[BP-2.2][FLINK-40904][FLINK-40780][runtime] Detect FORWARD parallelism mismatch correctly and make the downgrade configurable

docs/layouts/shortcodes/generated/pipeline_configuration.html

@@ -56,6 +56,12 @@
             <td>Boolean</td>
             <td>Forces Flink to register avro classes in kryo serializer.<br /><br />Important: Make sure to include the flink-avro module. Otherwise, nothing will be registered. For backward compatibility, the default value is empty to conform to the behavior of the older version. That is, always register avro with kryo, and if flink-avro is not in the class path, register a dummy serializer. In Flink-2.0, we will set the default value to true.</td>
         </tr>
+        <tr>
+            <td><h5>pipeline.forward-edge.parallelism-mismatch-mode</h5></td>
+            <td style="word-wrap: break-word;">REBALANCE</td>
+            <td><p>Enum</p></td>
+            <td>Determines how the runtime handles a FORWARD (pointwise) edge whose producer and consumer parallelism no longer match, which can happen when the parallelism of connected operators is changed independently (for example via the AdaptiveScheduler or <code class="highlighter-rouge">pipeline.jobvertex-parallelism-overrides</code>). A FORWARD edge is only valid at equal parallelism; on a mismatch one of the following strategies is applied:<ul><li><code class="highlighter-rouge">REBALANCE</code>: replace the forward partitioner with a rebalance (round-robin) partitioner. This preserves throughput but reorders records, which is safe for append-only streams but corrupts order-sensitive (changelog) streams.</li><li><code class="highlighter-rouge">FAIL</code>: fail the job with an exception instead of silently changing the data distribution.</li><li><code class="highlighter-rouge">KEEP_FORWARD</code>: keep the forward partitioner. Record order is preserved but the records are funneled to a single consumer subtask, so the remaining consumer subtasks stay idle.</li></ul><br /><br />Possible values:<ul><li>"REBALANCE"</li><li>"FAIL"</li><li>"KEEP_FORWARD"</li></ul></td>
+        </tr>
         <tr>
             <td><h5>pipeline.generic-types</h5></td>
             <td style="word-wrap: break-word;">true</td>

flink-core/src/main/java/org/apache/flink/configuration/PipelineOptions.java

@@ -304,6 +304,68 @@ public enum VertexDescriptionMode {
         CASCADING
     }
 
+    public static final ConfigOption<ForwardEdgeParallelismMismatchMode>
+            FORWARD_EDGE_PARALLELISM_MISMATCH_MODE =
+                    key("pipeline.forward-edge.parallelism-mismatch-mode")
+                            .enumType(ForwardEdgeParallelismMismatchMode.class)
+                            .defaultValue(ForwardEdgeParallelismMismatchMode.REBALANCE)
+                            .withDescription(
+                                    Description.builder()
+                                            .text(
+                                                    "Determines how the runtime handles a FORWARD (pointwise) edge whose"
+                                                            + " producer and consumer parallelism no longer match, which can happen"
+                                                            + " when the parallelism of connected operators is changed"
+                                                            + " independently (for example via the AdaptiveScheduler or %s). A"
+                                                            + " FORWARD edge is only valid at equal parallelism; on a mismatch one"
+                                                            + " of the following strategies is applied:",
+                                                    code(PARALLELISM_OVERRIDES.key()))
+                                            .list(
+                                                    text(
+                                                            "%s: replace the forward partitioner with a rebalance (round-robin)"
+                                                                    + " partitioner. This preserves throughput but reorders records,"
+                                                                    + " which is safe for append-only streams but corrupts"
+                                                                    + " order-sensitive (changelog) streams.",
+                                                            code(
+                                                                    ForwardEdgeParallelismMismatchMode
+                                                                            .REBALANCE
+                                                                            .name())),
+                                                    text(
+                                                            "%s: fail the job with an exception instead of silently changing"
+                                                                    + " the data distribution.",
+                                                            code(
+                                                                    ForwardEdgeParallelismMismatchMode
+                                                                            .FAIL
+                                                                            .name())),
+                                                    text(
+                                                            "%s: keep the forward partitioner. Record order is preserved but"
+                                                                    + " the records are funneled to a single consumer subtask, so the"
+                                                                    + " remaining consumer subtasks stay idle.",
+                                                            code(
+                                                                    ForwardEdgeParallelismMismatchMode
+                                                                            .KEEP_FORWARD
+                                                                            .name())))
+                                            .build());
+
+    /**
+     * The strategy applied when a FORWARD (pointwise) edge connects a producer and consumer with
+     * mismatched parallelism.
+     */
+    @PublicEvolving
+    public enum ForwardEdgeParallelismMismatchMode {
+        /**
+         * Replace the forward partitioner with a rebalance partitioner. Reorders records; safe only
+         * for append-only streams.
+         */
+        REBALANCE,
+        /** Fail the job with an exception. */
+        FAIL,
+        /**
+         * Keep the forward partitioner. Preserves record order but funnels records to a single
+         * consumer subtask.
+         */
+        KEEP_FORWARD
+    }
+
     public static final ConfigOption<Boolean> VERTEX_NAME_INCLUDE_INDEX_PREFIX =
             key("pipeline.vertex-name-include-index-prefix")
                     .booleanType()

flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/runtime/SavepointEnvironment.java

@@ -314,6 +314,11 @@ public ResultPartitionWriter getWriter(int index) {
         throw new UnsupportedOperationException(ERROR_MSG);
     }
 
+    @Override
+    public int getWriterConsumerParallelism(int index) {
+        throw new UnsupportedOperationException(ERROR_MSG);
+    }
+
     @Override
     public ResultPartitionWriter[] getAllWriters() {
         return new ResultPartitionWriter[0];

flink-runtime/src/main/java/org/apache/flink/runtime/deployment/ResultPartitionDeploymentDescriptor.java

@@ -28,6 +28,7 @@
 
 import java.io.Serializable;
 
+import static org.apache.flink.util.Preconditions.checkArgument;
 import static org.apache.flink.util.Preconditions.checkNotNull;
 
 /**
@@ -45,14 +46,31 @@ public class ResultPartitionDeploymentDescriptor implements Serializable {
 
     private final int maxParallelism;
 
+    /** Parallelism of the consuming job vertices, or {@link #UNKNOWN_CONSUMER_PARALLELISM}. */
+    private final int consumerParallelism;
+
+    /**
+     * The consumer parallelism is not decided yet when the producer is deployed. This only happens
+     * in dynamic graphs (AdaptiveBatchScheduler), where the parallelism of a consumer vertex may be
+     * decided after its producers are deployed. Every other partition carries the actual consumer
+     * parallelism.
+     */
+    public static final int UNKNOWN_CONSUMER_PARALLELISM = -1;
+
     public ResultPartitionDeploymentDescriptor(
             PartitionDescriptor partitionDescriptor,
             ShuffleDescriptor shuffleDescriptor,
-            int maxParallelism) {
+            int maxParallelism,
+            int consumerParallelism) {
         this.partitionDescriptor = checkNotNull(partitionDescriptor);
         this.shuffleDescriptor = checkNotNull(shuffleDescriptor);
         KeyGroupRangeAssignment.checkParallelismPreconditions(maxParallelism);
         this.maxParallelism = maxParallelism;
+        checkArgument(
+                consumerParallelism > 0 || consumerParallelism == UNKNOWN_CONSUMER_PARALLELISM,
+                "Invalid consumer parallelism %s.",
+                consumerParallelism);
+        this.consumerParallelism = consumerParallelism;
     }
 
     public IntermediateDataSetID getResultId() {
@@ -88,6 +106,16 @@ public int getMaxParallelism() {
         return maxParallelism;
     }
 
+    /**
+     * Returns the parallelism of the job vertices consuming this partition, or {@link
+     * #UNKNOWN_CONSUMER_PARALLELISM} if it is not decided yet. Unlike {@link
+     * #getNumberOfSubpartitions()}, this is the actual consumer parallelism for every distribution
+     * pattern.
+     */
+    public int getConsumerParallelism() {
+        return consumerParallelism;
+    }
+
     public ShuffleDescriptor getShuffleDescriptor() {
         return shuffleDescriptor;
     }

flink-runtime/src/main/java/org/apache/flink/runtime/execution/Environment.java

@@ -31,6 +31,7 @@
 import org.apache.flink.runtime.checkpoint.CheckpointMetrics;
 import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
 import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriteRequestExecutorFactory;
+import org.apache.flink.runtime.deployment.ResultPartitionDeploymentDescriptor;
 import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
 import org.apache.flink.runtime.externalresource.ExternalResourceInfoProvider;
 import org.apache.flink.runtime.io.disk.iomanager.IOManager;
@@ -247,6 +248,13 @@ void acknowledgeCheckpoint(
 
     ResultPartitionWriter getWriter(int index);
 
+    /**
+     * Returns the parallelism of the job vertices consuming the partition written by {@link
+     * #getWriter(int)}, or {@link ResultPartitionDeploymentDescriptor#UNKNOWN_CONSUMER_PARALLELISM}
+     * if it is not decided yet.
+     */
+    int getWriterConsumerParallelism(int index);
+
     ResultPartitionWriter[] getAllWriters();
 
     IndexedInputGate getInputGate(int index);

flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java

@@ -565,7 +565,10 @@ private static ResultPartitionDeploymentDescriptor createResultPartitionDeployme
             IntermediateResultPartition partition,
             ShuffleDescriptor shuffleDescriptor) {
         return new ResultPartitionDeploymentDescriptor(
-                partitionDescriptor, shuffleDescriptor, getPartitionMaxParallelism(partition));
+                partitionDescriptor,
+                shuffleDescriptor,
+                getPartitionMaxParallelism(partition),
+                partition.getIntermediateResult().getConsumersParallelism());
     }
 
     /**

flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/IntermediateDataSet.java

@@ -146,6 +146,12 @@ public void updateOutputPattern(
         this.isForward = isForward;
     }
 
+    public void updateDistributionPattern(
+            boolean isForward, DistributionPattern distributionPattern) {
+        this.isForward = isForward;
+        this.distributionPattern = distributionPattern;
+    }
+
     public void increaseNumJobEdgesToCreate() {
         this.numJobEdgesToCreate++;
     }

flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobEdge.java

@@ -35,7 +35,7 @@ public class JobEdge implements java.io.Serializable {
     private final JobVertex target;
 
     /** The distribution pattern that should be used for this job edge. */
-    private final DistributionPattern distributionPattern;
+    private DistributionPattern distributionPattern;
 
     /** The channel rescaler that should be used for this job edge on downstream side. */
     private SubtaskStateMapper downstreamSubtaskStateMapper = SubtaskStateMapper.ROUND_ROBIN;
@@ -54,7 +54,7 @@ public class JobEdge implements java.io.Serializable {
 
     private final boolean isBroadcast;
 
-    private final boolean isForward;
+    private boolean isForward;
 
     /**
      * Optional name for the pre-processing operation (sort, combining sort, ...), to be displayed
@@ -111,6 +111,20 @@ public JobEdge(
         this.intraInputKeyCorrelated = intraInputKeyCorrelated;
     }
 
+    public void updateDistributionPattern(
+            boolean isForward,
+            DistributionPattern distributionPattern,
+            SubtaskStateMapper upstreamSubtaskStateMapper,
+            SubtaskStateMapper downstreamSubtaskStateMapper,
+            String shipStrategyName) {
+        this.isForward = isForward;
+        this.distributionPattern = distributionPattern;
+        this.upstreamSubtaskStateMapper = checkNotNull(upstreamSubtaskStateMapper);
+        this.downstreamSubtaskStateMapper = checkNotNull(downstreamSubtaskStateMapper);
+        this.shipStrategyName = shipStrategyName;
+        this.source.updateDistributionPattern(isForward, distributionPattern);
+    }
+
     /**
      * Returns the data set at the source of the edge. May be null, if the edge refers to the source
      * via an ID and has not been connected.

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveScheduler.java

@@ -160,6 +160,7 @@
 import static org.apache.flink.configuration.JobManagerOptions.SCHEDULER_RESCALE_TRIGGER_MAX_DELAY;
 import static org.apache.flink.configuration.TraceOptions.CHECKPOINT_SPAN_DETAIL_LEVEL;
 import static org.apache.flink.runtime.executiongraph.ExecutionGraphUtils.isAnyOutputBlocking;
+import static org.apache.flink.runtime.scheduler.adaptive.ForwardEdgesAdapter.copyJobGraphWithAdaptedForwardEdges;
 
 /**
  * A {@link SchedulerNG} implementation that uses the declarative resource management and
@@ -1473,8 +1474,11 @@ private ExecutionGraph createExecutionGraphAndRestoreState(
                                     new ExecutionStateUpdateListener[0]));
         }
 
+        final JobGraph adjustedJobGraph =
+                copyJobGraphWithAdaptedForwardEdges(jobInformation, adjustedParallelismStore);
+
         return executionGraphFactory.createAndRestoreExecutionGraph(
-                jobInformation.copyJobGraph(),
+                adjustedJobGraph,
                 completedCheckpointStore,
                 checkpointsCleaner,
                 checkpointIdCounter,

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/ForwardEdgesAdapter.java

@@ -0,0 +1,75 @@
+/*
+ * 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.runtime.scheduler.adaptive;
+
+import org.apache.flink.runtime.jobgraph.JobEdge;
+import org.apache.flink.runtime.jobgraph.JobGraph;
+import org.apache.flink.runtime.jobgraph.JobVertex;
+import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.scheduler.VertexParallelismStore;
+import org.apache.flink.streaming.runtime.partitioner.ForwardPartitioner;
+import org.apache.flink.streaming.runtime.partitioner.RebalancePartitioner;
+import org.apache.flink.streaming.runtime.partitioner.StreamPartitioner;
+import org.apache.flink.util.Preconditions;
+
+import static org.apache.flink.runtime.jobgraph.DistributionPattern.ALL_TO_ALL;
+import static org.apache.flink.runtime.jobgraph.DistributionPattern.POINTWISE;
+
+/** The utils to try automatic adapt the setting of forward type job edges. */
+class ForwardEdgesAdapter {
+
+    private ForwardEdgesAdapter() {}
+
+    /**
+     * Get the copied job graph that was completed the switching between FORWARD and REBALANCE for
+     * {@link AdaptiveScheduler}.
+     */
+    static JobGraph copyJobGraphWithAdaptedForwardEdges(
+            JobGraphJobInformation jobInformation, VertexParallelismStore vertexParallelismStore) {
+        final JobGraph adjustedGraph = jobInformation.copyJobGraph();
+        for (JobVertex vertex : adjustedGraph.getVerticesSortedTopologicallyFromSources()) {
+            for (JobEdge input : vertex.getInputs()) {
+                if (input.isForward()) {
+                    Preconditions.checkState(
+                            input.getSource().isForward(),
+                            "Unexpected chained result occurred to here, may be a bug.");
+                    final boolean forwardable = isSameParallelism(input, vertexParallelismStore);
+                    StreamPartitioner<?> partitioner =
+                            forwardable ? new ForwardPartitioner<>() : new RebalancePartitioner<>();
+                    input.updateDistributionPattern(
+                            forwardable,
+                            forwardable ? POINTWISE : ALL_TO_ALL,
+                            partitioner.getUpstreamSubtaskStateMapper(),
+                            partitioner.getDownstreamSubtaskStateMapper(),
+                            partitioner.toString());
+                }
+            }
+        }
+        return adjustedGraph;
+    }
+
+    static boolean isSameParallelism(JobEdge jobEdge, VertexParallelismStore vertexParallelism) {
+        JobVertexID upstreamVertexID = jobEdge.getSource().getProducer().getID();
+        JobVertexID downstreamVertexID = jobEdge.getTarget().getID();
+        int upParallelism = vertexParallelism.getParallelismInfo(upstreamVertexID).getParallelism();
+        int downParallelism =
+                vertexParallelism.getParallelismInfo(downstreamVertexID).getParallelism();
+        return upParallelism == downParallelism;
+    }
+}

flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java

@@ -57,6 +57,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 
+import static org.apache.flink.util.Preconditions.checkArgument;
 import static org.apache.flink.util.Preconditions.checkNotNull;
 import static org.apache.flink.util.Preconditions.checkState;
 
@@ -92,6 +93,7 @@ public class RuntimeEnvironment implements Environment {
     private final Map<String, Future<Path>> distCacheEntries;
 
     private final ResultPartitionWriter[] writers;
+    private final int[] writerConsumerParallelisms;
     private final IndexedInputGate[] inputGates;
 
     private final TaskEventDispatcher taskEventDispatcher;
@@ -142,6 +144,7 @@ public RuntimeEnvironment(
             InputSplitProvider splitProvider,
             Map<String, Future<Path>> distCacheEntries,
             ResultPartitionWriter[] writers,
+            int[] writerConsumerParallelisms,
             IndexedInputGate[] inputGates,
             TaskEventDispatcher taskEventDispatcher,
             CheckpointResponder checkpointResponder,
@@ -174,6 +177,8 @@ public RuntimeEnvironment(
         this.splitProvider = checkNotNull(splitProvider);
         this.distCacheEntries = checkNotNull(distCacheEntries);
         this.writers = checkNotNull(writers);
+        this.writerConsumerParallelisms = checkNotNull(writerConsumerParallelisms);
+        checkArgument(writers.length == writerConsumerParallelisms.length);
         this.inputGates = checkNotNull(inputGates);
         this.taskEventDispatcher = checkNotNull(taskEventDispatcher);
         this.checkpointResponder = checkNotNull(checkpointResponder);
@@ -303,6 +308,11 @@ public ResultPartitionWriter getWriter(int index) {
         return writers[index];
     }
 
+    @Override
+    public int getWriterConsumerParallelism(int index) {
+        return writerConsumerParallelisms[index];
+    }
+
     @Override
     public ResultPartitionWriter[] getAllWriters() {
         return writers;