From dd28e3842c776bccc44077951fef6be68ca5e543 Mon Sep 17 00:00:00 2001 From: Rion Williams Date: Mon, 13 Jul 2026 17:27:39 -0500 Subject: [PATCH] [FLINK-40113][checkpoint] Reject restore when operators on a keyed vertex disagree on max parallelism During restore, each operator's recorded max parallelism is reconciled onto its shared vertex individually. When chaining has regrouped operators since the checkpoint, they can disagree, and the reconciliation silently keeps whichever value is applied last -- restoring a keyed operator under a foreign key-group count. Before the loop, verify that all operators chained into a vertex carrying keyed state agree, and fail with a clear error otherwise. The check spans all operators, since a non-keyed one's value can win the reconciliation and misroute a keyed operator's state. Adds ChainingMaxParallelismStateLossITCase covering both directions (explicit value below and above the chain head) on the HashMap and RocksDB backends. --- .../checkpoint/StateAssignmentOperation.java | 67 ++++ .../StateAssignmentOperationTest.java | 72 ++++ ...ChainingMaxParallelismStateLossITCase.java | 363 ++++++++++++++++++ 3 files changed, 502 insertions(+) create mode 100644 flink-tests/src/test/java/org/apache/flink/test/checkpointing/ChainingMaxParallelismStateLossITCase.java diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperation.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperation.java index 6bc30d488c4f2b..eaa624be10b6f6 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperation.java +++ b/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. + * + *

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 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 keyGroupPartitions, TaskStateAssignment stateAssignment) { stateAssignment.oldState.forEach( diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperationTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperationTest.java index 711c3ef5cf385e..1f2c9c7a3031c1 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperationTest.java +++ b/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 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 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. diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/ChainingMaxParallelismStateLossITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/ChainingMaxParallelismStateLossITCase.java new file mode 100644 index 00000000000000..3f7d47737bc488 --- /dev/null +++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/ChainingMaxParallelismStateLossITCase.java @@ -0,0 +1,363 @@ +/* + * 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.test.checkpointing; + +import org.apache.flink.api.common.JobID; +import org.apache.flink.api.common.functions.OpenContext; +import org.apache.flink.api.common.functions.RichMapFunction; +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDescriptor; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDescriptor; +import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.client.program.ClusterClient; +import org.apache.flink.configuration.CheckpointingOptions; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.PipelineOptions; +import org.apache.flink.configuration.StateBackendOptions; +import org.apache.flink.core.execution.SavepointFormatType; +import org.apache.flink.runtime.jobgraph.JobGraph; +import org.apache.flink.runtime.jobgraph.SavepointRestoreSettings; +import org.apache.flink.runtime.minicluster.MiniCluster; +import org.apache.flink.runtime.state.FunctionInitializationContext; +import org.apache.flink.runtime.state.FunctionSnapshotContext; +import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration; +import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction; +import org.apache.flink.streaming.api.datastream.DataStream; +import org.apache.flink.streaming.api.datastream.KeyedStream; +import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.api.functions.KeyedProcessFunction; +import org.apache.flink.streaming.api.functions.sink.legacy.SinkFunction; +import org.apache.flink.streaming.api.functions.source.legacy.RichParallelSourceFunction; +import org.apache.flink.streaming.util.RestartStrategyUtils; +import org.apache.flink.test.junit5.InjectClusterClient; +import org.apache.flink.test.junit5.InjectMiniCluster; +import org.apache.flink.test.junit5.MiniClusterExtension; +import org.apache.flink.util.Collector; + +import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import java.nio.file.Path; +import java.time.Duration; +import java.util.Collections; +import java.util.Map; +import java.util.TreeMap; +import java.util.concurrent.ConcurrentHashMap; + +import static org.apache.flink.runtime.testutils.CommonTestUtils.waitForAllTaskRunning; +import static org.apache.flink.test.util.TestUtils.submitJobAndWaitForResult; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Verifies that restoring a savepoint is rejected when a chaining change places operators with + * different recorded max parallelism onto a single keyed vertex. + * + *

A keyed operator is a normal {@code keyBy} chain head with an auto-derived max parallelism + * ({@value #CHAIN_HEAD_MAX_PARALLELISM} for parallelism 1). A downstream stateful operator carrying + * an explicit, different max parallelism chains after it. With chaining OFF the two are separate + * vertices and each records its own value; with chaining ON they merge into one keyed vertex that + * now carries two different recorded values. The savepoint's key-group count for the keyed operator + * therefore cannot be reconciled with the downstream operator's, so restore is rejected with a + * clear error instead of remapping keyed state through an incompatible key-group layout -- the same + * outcome whether the downstream value is below or above the chain head's. + */ +class ChainingMaxParallelismStateLossITCase { + + private static final int NUM_KEYS = 4; + private static final long JOB1_PER_KEY = 100; + private static final long JOB2_PER_KEY = 50; + + /** Auto-derived max parallelism of the keyed operator (chain head) at parallelism 1. */ + private static final int CHAIN_HEAD_MAX_PARALLELISM = 128; + + private static final int EXPLICIT_BELOW_HEAD = 64; + private static final int EXPLICIT_ABOVE_HEAD = 256; + + /** Final running count observed per key (per-key counts are monotonic). */ + private static final Map COUNTS = new ConcurrentHashMap<>(); + + @RegisterExtension + private static final MiniClusterExtension MINI_CLUSTER = + new MiniClusterExtension( + new MiniClusterResourceConfiguration.Builder() + .setNumberTaskManagers(1) + .setNumberSlotsPerTaskManager(4) + .build()); + + @TempDir private Path tempDir; + + @ParameterizedTest(name = "backend={0}") + @ValueSource(strings = {"hashmap", "rocksdb"}) + void rejectsRestoreWhenDownstreamMaxParallelismBelowChainHead( + String backend, + @InjectClusterClient ClusterClient client, + @InjectMiniCluster MiniCluster miniCluster) + throws Exception { + assertThatThrownBy( + () -> + savepointChainedOffRestoreChainedOn( + EXPLICIT_BELOW_HEAD, backend, client, miniCluster)) + .hasStackTraceContaining("recorded different maximum parallelism"); + } + + @ParameterizedTest(name = "backend={0}") + @ValueSource(strings = {"hashmap", "rocksdb"}) + void rejectsRestoreWhenDownstreamMaxParallelismAboveChainHead( + String backend, + @InjectClusterClient ClusterClient client, + @InjectMiniCluster MiniCluster miniCluster) + throws Exception { + assertThatThrownBy( + () -> + savepointChainedOffRestoreChainedOn( + EXPLICIT_ABOVE_HEAD, backend, client, miniCluster)) + .hasStackTraceContaining("recorded different maximum parallelism"); + } + + /** + * Runs one savepoint (chaining OFF, keyed operator and downstream operator on separate vertices + * at their own max parallelism) then restore (chaining ON, downstream operator chained under + * the keyed head) cycle, returning the per-key counts if the restore is not rejected. + */ + private Map savepointChainedOffRestoreChainedOn( + int downstreamMaxParallelism, + String backend, + ClusterClient client, + MiniCluster miniCluster) + throws Exception { + COUNTS.clear(); + + // Job 1 (chaining OFF): drive each key to JOB1_PER_KEY, then savepoint and cancel. + final JobGraph job1 = + buildJobGraph(false, JOB1_PER_KEY, false, downstreamMaxParallelism, backend); + final JobID jobId1 = job1.getJobID(); + client.submitJob(job1).get(); + waitForAllTaskRunning(miniCluster, jobId1, false); + waitUntilAllKeysReach(JOB1_PER_KEY); + + final String savepoint = + client.triggerSavepoint( + jobId1, + tempDir.resolve("savepoints").toUri().toString(), + SavepointFormatType.CANONICAL) + .get(); + client.cancel(jobId1).get(); + waitUntilNoJobRunning(client); + + // Job 2 (chaining ON): the downstream operator chains under the keyed head, whose max + // parallelism differs from the downstream operator's explicit one, so restore is rejected. + COUNTS.clear(); + final JobGraph job2 = + buildJobGraph(true, JOB2_PER_KEY, true, downstreamMaxParallelism, backend); + job2.setSavepointRestoreSettings(SavepointRestoreSettings.forPath(savepoint)); + submitJobAndWaitForResult(client, job2, getClass().getClassLoader()); + + return new TreeMap<>(COUNTS); + } + + private JobGraph buildJobGraph( + boolean chaining, + long elementsPerKey, + boolean terminate, + int downstreamMaxParallelism, + String backend) { + // State backend and checkpoint storage are configured per job so a single shared cluster + // can run both backends across the parameterized cases. + final Configuration config = new Configuration(); + config.set(StateBackendOptions.STATE_BACKEND, backend); + config.set( + CheckpointingOptions.CHECKPOINTS_DIRECTORY, + tempDir.resolve("checkpoints").toUri().toString()); + // Default is already true; set explicitly for clarity -- this is what lets the downstream + // operator chain under a keyed head with a different (auto-derived) max parallelism. + config.set( + PipelineOptions.OPERATOR_CHAINING_CHAIN_OPERATORS_WITH_DIFFERENT_MAX_PARALLELISM, + true); + + final StreamExecutionEnvironment env = + StreamExecutionEnvironment.getExecutionEnvironment(config); + env.setParallelism(1); + // No env-level max parallelism, so the keyed (chain-head) operator uses an auto-derived + // value + // while the downstream operator carries its own explicit one. + env.enableCheckpointing(Duration.ofMinutes(10).toMillis()); + RestartStrategyUtils.configureNoRestartStrategy(env); + if (!chaining) { + env.disableOperatorChaining(); + } + + final DataStream source = + env.addSource(new ControllableSource(NUM_KEYS, elementsPerKey, terminate)) + .uid("src") + .name("src"); + + // A plain keyBy makes the keyed operator a chain head (the keyBy hash edge is a chain + // break). + final KeyedStream keyed = source.keyBy(value -> value % NUM_KEYS); + + final SingleOutputStreamOperator> counted = + keyed.process(new PerKeyCounter()).name("keyed").uid("keyed"); + + // A forward (chainable) edge to a downstream stateful operator with an explicit, different + // max parallelism: it becomes a chained non-head under the keyed head when chaining is on. + final SingleOutputStreamOperator> mapped = + counted.map(new StatefulPassThrough()) + .name("mapped") + .uid("mapped") + .setMaxParallelism(downstreamMaxParallelism); + + mapped.addSink(new CountsCollectingSink()).uid("sink").name("sink"); + + return env.getStreamGraph().getJobGraph(); + } + + private void waitUntilAllKeysReach(long target) throws InterruptedException { + while (true) { + final Map current = new TreeMap<>(COUNTS); + if (current.size() == NUM_KEYS + && current.values().stream().allMatch(v -> v >= target)) { + return; + } + Thread.sleep(25); + } + } + + private void waitUntilNoJobRunning(ClusterClient client) throws Exception { + while (!client.listJobs().get().stream() + .allMatch(s -> s.getJobState().isGloballyTerminalState())) { + Thread.sleep(50); + } + } + + /** + * Emits each of {@code numKeys} keys {@code elementsPerKey} times, then either terminates or + * stays alive (sleeping) so a savepoint can be taken while the job runs. + */ + private static final class ControllableSource extends RichParallelSourceFunction { + + private static final long serialVersionUID = 1L; + + private final int numKeys; + private final long elementsPerKey; + private final boolean terminateAfterEmission; + + private volatile boolean running = true; + + ControllableSource(int numKeys, long elementsPerKey, boolean terminateAfterEmission) { + this.numKeys = numKeys; + this.elementsPerKey = elementsPerKey; + this.terminateAfterEmission = terminateAfterEmission; + } + + @Override + public void run(SourceContext ctx) throws Exception { + final Object lock = ctx.getCheckpointLock(); + for (long i = 0; i < elementsPerKey && running; i++) { + synchronized (lock) { + for (int key = 0; key < numKeys; key++) { + ctx.collect(key); + } + } + } + if (terminateAfterEmission) { + return; + } + // Stay alive (without emitting more) so the state is frozen while a savepoint is taken. + while (running) { + Thread.sleep(50); + } + } + + @Override + public void cancel() { + running = false; + } + } + + /** Per-key monotonic counter backed by keyed {@link ValueState}. */ + private static final class PerKeyCounter + extends KeyedProcessFunction> { + + private static final long serialVersionUID = 1L; + + private transient ValueState counter; + + @Override + public void open(OpenContext openContext) { + counter = + getRuntimeContext().getState(new ValueStateDescriptor<>("counter", Long.class)); + } + + @Override + public void processElement(Integer value, Context ctx, Collector> out) + throws Exception { + final Long previous = counter.value(); + final long next = (previous == null ? 0L : previous) + 1L; + counter.update(next); + out.collect(Tuple2.of(ctx.getCurrentKey(), next)); + } + } + + /** + * Pass-through map that keeps operator (non-keyed) list state, so it records an operator state + * with its vertex's max parallelism in the savepoint. + */ + private static final class StatefulPassThrough + extends RichMapFunction, Tuple2> + implements CheckpointedFunction { + + private static final long serialVersionUID = 1L; + + private transient ListState operatorState; + + @Override + public Tuple2 map(Tuple2 value) { + return value; + } + + @Override + public void snapshotState(FunctionSnapshotContext context) throws Exception { + operatorState.update(Collections.singletonList(1L)); + } + + @Override + public void initializeState(FunctionInitializationContext context) throws Exception { + operatorState = + context.getOperatorStateStore() + .getListState(new ListStateDescriptor<>("op", Long.class)); + } + } + + /** + * Records the maximum count seen per key so the test thread can read the final per-key counts. + */ + private static final class CountsCollectingSink implements SinkFunction> { + + private static final long serialVersionUID = 1L; + + @Override + public void invoke(Tuple2 value, Context context) { + COUNTS.merge(value.f0, value.f1, Math::max); + } + } +}