[FLINK-37084][python] Fix TimerRegistration concurrency issue in PyFlink (#26004)

Author: suez1224

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

@@ -44,6 +44,8 @@ public interface PythonFunctionRunner extends AutoCloseable {
     /** Send the triggered timer to the Python function. */
     void processTimer(byte[] timerData) throws Exception;
 
+    void drainUnregisteredTimers();
+
     /**
      * Retrieves the Python function result.
      *

flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractPythonFunctionOperator.java

@@ -268,17 +268,25 @@ public Configuration getConfiguration() {
 
     protected abstract PythonEnvironmentManager createPythonEnvironmentManager();
 
+    protected void drainUnregisteredTimers() {}
+
     /**
      * Advances the watermark of all managed timer services, potentially firing event time timers.
      * It also ensures that the fired timers are processed in the Python user-defined functions.
      */
     private void advanceWatermark(Watermark watermark) throws Exception {
         if (getTimeServiceManager().isPresent()) {
             InternalTimeServiceManager<?> timeServiceManager = getTimeServiceManager().get();
+            // make sure the registered timer are processed before advancing the watermark to
+            // ensure the timers could be triggered
+            drainUnregisteredTimers();
             timeServiceManager.advanceWatermark(watermark);
 
             while (!isBundleFinished()) {
                 invokeFinishBundle();
+                // make sure the registered timer are processed before advancing the watermark to
+                // ensure the timers could be triggered
+                drainUnregisteredTimers();
                 timeServiceManager.advanceWatermark(watermark);
             }
         }

flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/AbstractExternalPythonFunctionOperator.java

@@ -136,6 +136,11 @@ protected ProcessPythonEnvironmentManager createPythonEnvironmentManager() {
         }
     }
 
+    @Override
+    protected void drainUnregisteredTimers() {
+        pythonFunctionRunner.drainUnregisteredTimers();
+    }
+
     protected void emitResults() throws Exception {
         Tuple3<String, byte[], Integer> resultTuple;
         while ((resultTuple = pythonFunctionRunner.pollResult()) != null && resultTuple.f2 != 0) {

flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/ExternalPythonCoProcessOperator.java

@@ -72,6 +72,7 @@ public void open() throws Exception {
     @Override
     public PythonFunctionRunner createPythonFunctionRunner() throws Exception {
         return new BeamDataStreamPythonFunctionRunner(
+                getContainingTask().getEnvironment(),
                 getRuntimeContext().getTaskInfo().getTaskName(),
                 createPythonEnvironmentManager(),
                 STATELESS_FUNCTION_URN,

flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/ExternalPythonKeyedCoProcessOperator.java

@@ -113,6 +113,7 @@ public void open() throws Exception {
     @Override
     public PythonFunctionRunner createPythonFunctionRunner() throws Exception {
         return new BeamDataStreamPythonFunctionRunner(
+                getContainingTask().getEnvironment(),
                 getRuntimeContext().getTaskInfo().getTaskName(),
                 createPythonEnvironmentManager(),
                 STATEFUL_FUNCTION_URN,

flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/ExternalPythonKeyedProcessOperator.java

@@ -139,6 +139,7 @@ public void onProcessingTime(InternalTimer<Row, Object> timer) throws Exception
     @Override
     public PythonFunctionRunner createPythonFunctionRunner() throws Exception {
         return new BeamDataStreamPythonFunctionRunner(
+                getContainingTask().getEnvironment(),
                 getRuntimeContext().getTaskInfo().getTaskName(),
                 createPythonEnvironmentManager(),
                 STATEFUL_FUNCTION_URN,

flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/ExternalPythonProcessOperator.java

@@ -68,6 +68,7 @@ public void open() throws Exception {
     @Override
     public PythonFunctionRunner createPythonFunctionRunner() throws Exception {
         return new BeamDataStreamPythonFunctionRunner(
+                getContainingTask().getEnvironment(),
                 getRuntimeContext().getTaskInfo().getTaskName(),
                 createPythonEnvironmentManager(),
                 STATELESS_FUNCTION_URN,

flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/process/timer/TimerRegistrationAction.java

@@ -0,0 +1,42 @@
+/*
+ * 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.timer;
+
+public class TimerRegistrationAction {
+
+    private final TimerRegistration timerRegistration;
+
+    private final byte[] serializedTimerData;
+
+    private boolean isRegistered;
+
+    public TimerRegistrationAction(
+            TimerRegistration timerRegistration, byte[] serializedTimerData) {
+        this.timerRegistration = timerRegistration;
+        this.serializedTimerData = serializedTimerData;
+        this.isRegistered = false;
+    }
+
+    public void run() {
+        if (!isRegistered) {
+            timerRegistration.setTimer(serializedTimerData);
+            isRegistered = true;
+        }
+    }
+}

flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamDataStreamPythonFunctionRunner.java

@@ -24,6 +24,7 @@
 import org.apache.flink.python.env.process.ProcessPythonEnvironmentManager;
 import org.apache.flink.python.metric.process.FlinkMetricContainer;
 import org.apache.flink.python.util.ProtoUtils;
+import org.apache.flink.runtime.execution.Environment;
 import org.apache.flink.runtime.memory.MemoryManager;
 import org.apache.flink.runtime.state.KeyedStateBackend;
 import org.apache.flink.runtime.state.OperatorStateBackend;
@@ -69,6 +70,7 @@ public class BeamDataStreamPythonFunctionRunner extends BeamPythonFunctionRunner
     private final List<FlinkFnApi.UserDefinedDataStreamFunction> userDefinedDataStreamFunctions;
 
     public BeamDataStreamPythonFunctionRunner(
+            Environment environment,
             String taskName,
             ProcessPythonEnvironmentManager environmentManager,
             String headOperatorFunctionUrn,
@@ -86,6 +88,7 @@ public BeamDataStreamPythonFunctionRunner(
             @Nullable FlinkFnApi.CoderInfoDescriptor timerCoderDescriptor,
             Map<String, FlinkFnApi.CoderInfoDescriptor> sideOutputCoderDescriptors) {
         super(
+                environment,
                 taskName,
                 environmentManager,
                 flinkMetricContainer,

flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java

@@ -31,11 +31,13 @@
 import org.apache.flink.python.env.process.ProcessPythonEnvironment;
 import org.apache.flink.python.env.process.ProcessPythonEnvironmentManager;
 import org.apache.flink.python.metric.process.FlinkMetricContainer;
+import org.apache.flink.runtime.execution.Environment;
 import org.apache.flink.runtime.memory.MemoryManager;
 import org.apache.flink.runtime.memory.OpaqueMemoryResource;
 import org.apache.flink.runtime.state.KeyedStateBackend;
 import org.apache.flink.runtime.state.OperatorStateBackend;
 import org.apache.flink.streaming.api.operators.python.process.timer.TimerRegistration;
+import org.apache.flink.streaming.api.operators.python.process.timer.TimerRegistrationAction;
 import org.apache.flink.streaming.api.runners.python.beam.state.BeamStateRequestHandler;
 import org.apache.flink.util.Preconditions;
 import org.apache.flink.util.ShutdownHookUtil;
@@ -85,6 +87,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -190,7 +193,12 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
 
     private transient Thread shutdownHook;
 
+    private transient Environment environment;
+
+    private transient List<TimerRegistrationAction> unregisteredTimers;
+
     public BeamPythonFunctionRunner(
+            Environment environment,
             String taskName,
             ProcessPythonEnvironmentManager environmentManager,
             @Nullable FlinkMetricContainer flinkMetricContainer,
@@ -204,6 +212,7 @@ public BeamPythonFunctionRunner(
             FlinkFnApi.CoderInfoDescriptor inputCoderDescriptor,
             FlinkFnApi.CoderInfoDescriptor outputCoderDescriptor,
             Map<String, FlinkFnApi.CoderInfoDescriptor> sideOutputCoderDescriptors) {
+        this.environment = environment;
         this.taskName = Preconditions.checkNotNull(taskName);
         this.environmentManager = Preconditions.checkNotNull(environmentManager);
         this.flinkMetricContainer = flinkMetricContainer;
@@ -301,6 +310,8 @@ public void open(ReadableConfig config) throws Exception {
         shutdownHook =
                 ShutdownHookUtil.addShutdownHook(
                         this, BeamPythonFunctionRunner.class.getSimpleName(), LOG);
+
+        unregisteredTimers = new LinkedList<>();
     }
 
     @Override
@@ -339,6 +350,14 @@ public void process(byte[] data) throws Exception {
         mainInputReceiver.accept(WindowedValue.valueInGlobalWindow(data));
     }
 
+    @Override
+    public void drainUnregisteredTimers() {
+        for (TimerRegistrationAction timerRegistrationAction : unregisteredTimers) {
+            timerRegistrationAction.run();
+        }
+        unregisteredTimers.clear();
+    }
+
     @Override
     public void processTimer(byte[] timerData) throws Exception {
         if (timerInputReceiver == null) {
@@ -681,7 +700,15 @@ public void onCompleted(BeamFnApi.ProcessBundleResponse response) {
 
     private TimerReceiverFactory createTimerReceiverFactory() {
         BiConsumer<Timer<?>, TimerInternals.TimerData> timerDataConsumer =
-                (timer, timerData) -> timerRegistration.setTimer((byte[]) timer.getUserKey());
+                (timer, timerData) -> {
+                    TimerRegistrationAction timerRegistrationAction =
+                            new TimerRegistrationAction(
+                                    timerRegistration, (byte[]) timer.getUserKey());
+                    unregisteredTimers.add(timerRegistrationAction);
+                    environment
+                            .getMainMailboxExecutor()
+                            .execute(timerRegistrationAction::run, "PythonTimerRegistration");
+                };
         return new TimerReceiverFactory(stageBundleFactory, timerDataConsumer, null);
     }