add flag in FlinkPipelineOptions to allow draining for pipelines with…

… RequiresStableInput
twitter-forks · Oct 18, 2023 · 4ca8d3a · 4ca8d3a
1 parent 65b2dcf
commit 4ca8d3a
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 1 deletion.
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java
@@ -299,6 +299,13 @@ public interface FlinkPipelineOptions
 
   void setFlinkConfDir(String confDir);
 
+  @Description("Allow drain operation for flink pipelines that contain RequiresStableInput operator. Note that at time of draining," +
+          "the RequiresStableInput contract might be violated if there any processing related failures in the DoFn operator.")
+  @Default.Boolean(false)
+  Boolean getEnableStableInputDrain();
+
+  void setEnableStableInputDrain(Boolean enableStableInputDrain);
+
   static FlinkPipelineOptions defaults() {
     return PipelineOptionsFactory.as(FlinkPipelineOptions.class);
   }

diff --git a/.../main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java b/.../main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java
@@ -199,6 +199,8 @@ public class DoFnOperator<InputT, OutputT>
   /** If true, we must process elements only after a checkpoint is finished. */
   final boolean requiresStableInput;
 
+  final boolean enableStableInputDrain;
+
   final int numConcurrentCheckpoints;
 
   private final boolean usesOnWindowExpiration;
@@ -323,6 +325,8 @@ public DoFnOperator(
               + Math.max(0, flinkOptions.getMinPauseBetweenCheckpoints()));
     }
 
+    this.enableStableInputDrain = flinkOptions.getEnableStableInputDrain();
+
     this.numConcurrentCheckpoints = flinkOptions.getNumConcurrentCheckpoints();
 
     this.finishBundleBeforeCheckpointing = flinkOptions.getFinishBundleBeforeCheckpointing();
@@ -626,7 +630,7 @@ void flushData() throws Exception {
     while (bundleStarted) {
       invokeFinishBundle();
     }
-    if (requiresStableInput) {
+    if (requiresStableInput && enableStableInputDrain) {
       // Flush any buffered events here before draining the pipeline. Note that this is best-effort
       // and requiresStableInput contract might be violated in cases where buffer processing fails.
       bufferingDoFnRunner.checkpointCompleted(Long.MAX_VALUE);

diff --git a/...t/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java b/...t/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java
@@ -2020,6 +2020,7 @@ public void testExactlyOnceBufferingFlushDuringDrain() throws Exception {
     FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
     options.setMaxBundleSize(2L);
     options.setCheckpointingInterval(1L);
+    options.setEnableStableInputDrain(true);
 
     TupleTag<String> outputTag = new TupleTag<>("main-output");
     WindowedValue.ValueOnlyWindowedValueCoder<String> windowedValueCoder =

diff --git a/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html b/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html
@@ -52,6 +52,11 @@
   <td>Disable Beam metrics in Flink Runner</td>
   <td>Default: <code>false</code></td>
 </tr>
+<tr>
+  <td><code>enableStableInputDrain</code></td>
+  <td>Allow drain operation for flink pipelines that contain RequiresStableInput operator. Note that at time of draining,the RequiresStableInput contract might be violated if there any processing related failures in the DoFn operator.</td>
+  <td>Default: <code>false</code></td>
+</tr>
 <tr>
   <td><code>executionModeForBatch</code></td>
   <td>Flink mode for data exchange of batch pipelines. Reference {@link org.apache.flink.api.common.ExecutionMode}. Set this to BATCH_FORCED if pipelines get blocked, see https://issues.apache.org/jira/browse/FLINK-10672</td>

diff --git a/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html b/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html
@@ -52,6 +52,11 @@
   <td>Disable Beam metrics in Flink Runner</td>
   <td>Default: <code>false</code></td>
 </tr>
+<tr>
+  <td><code>enable_stable_input_drain</code></td>
+  <td>Allow drain operation for flink pipelines that contain RequiresStableInput operator. Note that at time of draining,the RequiresStableInput contract might be violated if there any processing related failures in the DoFn operator.</td>
+  <td>Default: <code>false</code></td>
+</tr>
 <tr>
   <td><code>execution_mode_for_batch</code></td>
   <td>Flink mode for data exchange of batch pipelines. Reference {@link org.apache.flink.api.common.ExecutionMode}. Set this to BATCH_FORCED if pipelines get blocked, see https://issues.apache.org/jira/browse/FLINK-10672</td>