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 6bc30d488c4f2..eaa624be10b6f 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-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 0000000000000..be43375a92808 --- /dev/null +++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/ChainingMaxParallelismStateLossITCase.java @@ -0,0 +1,322 @@ +/* + * 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.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDescriptor; +import org.apache.flink.api.common.time.Deadline; +import org.apache.flink.api.java.functions.KeySelector; +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.testutils.MiniClusterResourceConfiguration; +import org.apache.flink.streaming.api.datastream.DataStream; +import org.apache.flink.streaming.api.datastream.DataStreamUtils; +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.util.MiniClusterWithClientResource; +import org.apache.flink.testutils.junit.utils.TempDirUtils; +import org.apache.flink.util.Collector; + +import org.junit.jupiter.api.AfterEach; +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.Map; +import java.util.TreeMap; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeUnit; + +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. + * + *

The keyed operator is a chained non-head (reached via {@link + * DataStreamUtils#reinterpretAsKeyedStream}) carrying an explicit max parallelism. With chaining + * OFF it is its own vertex and keeps that value, so its state is written into that many key groups. + * With chaining ON it chains under the auto-max-parallelism source head, whose derived value + * ({@value #CHAIN_HEAD_MAX_PARALLELISM} for parallelism 1) differs from the operator's explicit + * one. The savepoint's key-group count therefore differs from the restore vertex'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 explicit 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 source (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<>(); + + private MiniClusterWithClientResource cluster; + + @TempDir private static Path temporaryFolder; + + @AfterEach + void tearDown() { + if (cluster != null) { + cluster.after(); + cluster = null; + } + } + + @ParameterizedTest(name = "backend={0}") + @ValueSource(strings = {"hashmap", "rocksdb"}) + void rejectsRestoreWhenExplicitMaxParallelismBelowChainHead(String backend) throws Exception { + startCluster(backend); + + assertThatThrownBy(() -> savepointChainedOffRestoreChainedOn(EXPLICIT_BELOW_HEAD)) + .hasStackTraceContaining("recorded different maximum parallelism"); + } + + @ParameterizedTest(name = "backend={0}") + @ValueSource(strings = {"hashmap", "rocksdb"}) + void rejectsRestoreWhenExplicitMaxParallelismAboveChainHead(String backend) throws Exception { + startCluster(backend); + + assertThatThrownBy(() -> savepointChainedOffRestoreChainedOn(EXPLICIT_ABOVE_HEAD)) + .hasStackTraceContaining("recorded different maximum parallelism"); + } + + /** + * Runs one savepoint (chaining OFF, keyed operator on its own vertex at its explicit max + * parallelism) then restore (chaining ON, keyed operator chained under the source head) cycle, + * returning the per-key counts if the restore is not rejected. + */ + private Map savepointChainedOffRestoreChainedOn(int keyedMaxParallelism) + throws Exception { + COUNTS.clear(); + + final Deadline deadline = Deadline.now().plus(Duration.ofMinutes(2)); + final ClusterClient client = cluster.getClusterClient(); + + // Job 1 (chaining OFF): drive each key to JOB1_PER_KEY, then savepoint and cancel. + final JobGraph job1 = buildJobGraph(false, JOB1_PER_KEY, false, keyedMaxParallelism); + final JobID jobId1 = job1.getJobID(); + client.submitJob(job1).get(); + waitForAllTaskRunning(cluster.getMiniCluster(), jobId1, false); + waitUntilAllKeysReach(JOB1_PER_KEY, deadline); + + final String savepoint = + client.triggerSavepoint(jobId1, null, SavepointFormatType.CANONICAL) + .get(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS); + client.cancel(jobId1).get(); + waitUntilNoJobRunning(client); + + // Job 2 (chaining ON): the keyed operator chains under the source head, whose max + // parallelism differs from the operator's explicit one, so restore is rejected. + COUNTS.clear(); + final JobGraph job2 = buildJobGraph(true, JOB2_PER_KEY, true, keyedMaxParallelism); + job2.setSavepointRestoreSettings(SavepointRestoreSettings.forPath(savepoint)); + submitJobAndWaitForResult(client, job2, getClass().getClassLoader()); + + return new TreeMap<>(COUNTS); + } + + private JobGraph buildJobGraph( + boolean chaining, long elementsPerKey, boolean terminate, int keyedMaxParallelism) { + final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(1); + // No env-level max parallelism, so the (chain-head) source uses an auto-derived value while + // the keyed 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"); + + // reinterpretAsKeyedStream puts a forward (chainable) edge before the keyed operator, so it + // can become a chained non-head under the source head. + final KeyedStream keyed = + DataStreamUtils.reinterpretAsKeyedStream( + source, (KeySelector) value -> value % NUM_KEYS); + + final SingleOutputStreamOperator> counted = + keyed.process(new PerKeyCounter()) + .name("keyed") + .uid("keyed") + .setMaxParallelism(keyedMaxParallelism); + + counted.addSink(new CountsCollectingSink()).uid("sink").name("sink"); + + return env.getStreamGraph().getJobGraph(); + } + + private void startCluster(String backend) throws Exception { + final Configuration config = new Configuration(); + config.set(StateBackendOptions.STATE_BACKEND, backend); + config.set( + CheckpointingOptions.CHECKPOINTS_DIRECTORY, + TempDirUtils.newFolder(temporaryFolder).toURI().toString()); + config.set( + CheckpointingOptions.SAVEPOINT_DIRECTORY, + TempDirUtils.newFolder(temporaryFolder).toURI().toString()); + // Default is already true; set explicitly for clarity — this is what lets the keyed + // operator + // chain under a head with a different (auto-derived) max parallelism. + config.set( + PipelineOptions.OPERATOR_CHAINING_CHAIN_OPERATORS_WITH_DIFFERENT_MAX_PARALLELISM, + true); + + cluster = + new MiniClusterWithClientResource( + new MiniClusterResourceConfiguration.Builder() + .setConfiguration(config) + .setNumberTaskManagers(1) + .setNumberSlotsPerTaskManager(4) + .build()); + cluster.before(); + } + + private void waitUntilAllKeysReach(long target, Deadline deadline) throws InterruptedException { + while (true) { + final Map current = new TreeMap<>(COUNTS); + if (current.size() == NUM_KEYS + && current.values().stream().allMatch(v -> v >= target)) { + return; + } + if (!deadline.hasTimeLeft()) { + throw new IllegalStateException( + "Timed out waiting for all keys to reach " + target + "; saw " + current); + } + 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)); + } + } + + /** + * 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); + } + } +}