[FLINK-40113][checkpoint] Reject restore when operators on a keyed vertex disagree on max parallelism

flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperation.java

@@ -304,11 +304,78 @@ private void assignNonFinishedStateToTask(
     }
 
     public void checkParallelismPreconditions(TaskStateAssignment taskStateAssignment) {
+        checkMaxParallelismAgreement(taskStateAssignment);
         for (OperatorState operatorState : taskStateAssignment.oldState.values()) {
             checkParallelismPreconditions(operatorState, taskStateAssignment.executionJobVertex);
         }
     }
 
+    /**
+     * Verifies that all operators chained into a single keyed vertex recorded the same maximum
+     * parallelism in the checkpoint.
+     *
+     * <p>The per-operator reconciliation below ({@link
+     * #checkParallelismPreconditions(OperatorState, ExecutionJobVertex)}) adopts each operator's
+     * recorded maximum parallelism onto the shared vertex, so when operators disagree the vertex is
+     * left with whichever value is reconciled last. Any keyed operator on the vertex is then
+     * restored under that value rather than its own, remapping its state through an incompatible
+     * {@code hash % maxParallelism} layout. Operators sharing a vertex normally record its single
+     * maximum parallelism and therefore agree; they can only differ here if the chaining topology
+     * regrouped them since the checkpoint. This regrouping was permitted for graph construction but
+     * never validated on restore. A disagreeing operator need not be keyed itself: its recorded
+     * value can win the reconciliation and misroute another operator's keyed state, so all
+     * operators are compared. Vertices without keyed state are unaffected, since maximum
+     * parallelism only governs keyed-state routing.
+     */
+    private static void checkMaxParallelismAgreement(TaskStateAssignment taskStateAssignment) {
+        OperatorID referenceOperator = null;
+        int referenceMaxParallelism = -1;
+        OperatorID conflictingOperator = null;
+        int conflictingMaxParallelism = -1;
+        boolean vertexHasKeyedState = false;
+
+        for (Map.Entry<OperatorID, OperatorState> entry : taskStateAssignment.oldState.entrySet()) {
+            final OperatorState operatorState = entry.getValue();
+            vertexHasKeyedState |= hasKeyedState(operatorState);
+
+            if (referenceOperator == null) {
+                referenceOperator = entry.getKey();
+                referenceMaxParallelism = operatorState.getMaxParallelism();
+            } else if (conflictingOperator == null
+                    && operatorState.getMaxParallelism() != referenceMaxParallelism) {
+                conflictingOperator = entry.getKey();
+                conflictingMaxParallelism = operatorState.getMaxParallelism();
+            }
+        }
+
+        if (vertexHasKeyedState && conflictingOperator != null) {
+            throw new IllegalStateException(
+                    "The state for the execution job vertex "
+                            + taskStateAssignment.executionJobVertex.getJobVertexId()
+                            + " can not be restored. Operators "
+                            + referenceOperator
+                            + " and "
+                            + conflictingOperator
+                            + " are chained into the same keyed vertex but recorded different"
+                            + " maximum parallelism in the checkpoint ("
+                            + referenceMaxParallelism
+                            + " and "
+                            + conflictingMaxParallelism
+                            + "). Restoring would remap keyed state through an incompatible"
+                            + " key-group layout. This is currently not supported.");
+        }
+    }
+
+    private static boolean hasKeyedState(OperatorState operatorState) {
+        for (OperatorSubtaskState subtaskState : operatorState.getStates()) {
+            if (!subtaskState.getManagedKeyedState().isEmpty()
+                    || !subtaskState.getRawKeyedState().isEmpty()) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     private void reDistributeKeyedStates(
             List<KeyGroupRange> keyGroupPartitions, TaskStateAssignment stateAssignment) {
         stateAssignment.oldState.forEach(

flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperationTest.java

@@ -66,6 +66,7 @@
 import java.util.EnumMap;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Random;
@@ -96,6 +97,7 @@
 import static org.apache.flink.runtime.util.JobVertexConnectionUtils.connectNewDataSetAsInput;
 import static org.apache.flink.util.Preconditions.checkArgument;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /** Tests to verify state assignment operation. */
 class StateAssignmentOperationTest {
@@ -1315,6 +1317,76 @@ private List<OperatorID> buildOperatorIds(int numOperators) {
                 .collect(Collectors.toList());
     }
 
+    /**
+     * A keyed vertex whose chained operators recorded different maximum parallelism cannot be
+     * restored. The rejection must not depend on the order in which the operator states are
+     * reconciled onto the shared vertex, so both orders are exercised.
+     */
+    @ParameterizedTest
+    @ValueSource(booleans = {true, false})
+    void restoreRejectsKeyedVertexWithConflictingMaxParallelism(boolean keyedStateFirst)
+            throws Exception {
+        final OperatorID keyedOperator = new OperatorID();
+        final OperatorID chainedOperator = new OperatorID();
+        final int keyedMaxParallelism = 128;
+        final int chainedMaxParallelism = 64;
+
+        OperatorState keyedState =
+                new OperatorState(null, null, keyedOperator, 1, keyedMaxParallelism);
+        keyedState.putState(
+                0,
+                OperatorSubtaskState.builder()
+                        .setManagedKeyedState(
+                                StateObjectCollection.singleton(
+                                        createNewKeyedStateHandle(
+                                                KeyGroupRange.of(0, keyedMaxParallelism - 1))))
+                        .build());
+        OperatorState chainedState =
+                new OperatorState(null, null, chainedOperator, 1, chainedMaxParallelism);
+        chainedState.putState(
+                0,
+                OperatorSubtaskState.builder()
+                        .setManagedOperatorState(
+                                StateObjectCollection.singleton(
+                                        createNewOperatorStateHandle(2, new Random())))
+                        .build());
+
+        JobVertex jobVertex =
+                new JobVertex(
+                        "keyed-chain",
+                        new JobVertexID(),
+                        asList(
+                                OperatorIDPair.generatedIDOnly(keyedOperator),
+                                OperatorIDPair.generatedIDOnly(chainedOperator)));
+        jobVertex.setInvokableClass(NoOpInvokable.class);
+        jobVertex.setParallelism(1);
+        ExecutionJobVertex executionJobVertex =
+                ExecutionGraphTestUtils.getExecutionJobVertex(jobVertex);
+
+        Map<OperatorID, OperatorState> oldState = new LinkedHashMap<>();
+        if (keyedStateFirst) {
+            oldState.put(keyedOperator, keyedState);
+            oldState.put(chainedOperator, chainedState);
+        } else {
+            oldState.put(chainedOperator, chainedState);
+            oldState.put(keyedOperator, keyedState);
+        }
+
+        TaskStateAssignment taskStateAssignment =
+                new TaskStateAssignment(
+                        executionJobVertex, oldState, new HashMap<>(), new HashMap<>(), false);
+        StateAssignmentOperation stateAssignmentOperation =
+                new StateAssignmentOperation(
+                        1L, Collections.singleton(executionJobVertex), oldState, false, false);
+
+        assertThatThrownBy(
+                        () ->
+                                stateAssignmentOperation.checkParallelismPreconditions(
+                                        taskStateAssignment))
+                .isInstanceOf(IllegalStateException.class)
+                .hasMessageContaining("recorded different maximum parallelism");
+    }
+
     /**
      * Asserts the upstream output buffer state for a specific subtask by verifying the expected
      * upstream subtask and subpartition mappings.