rionmonster commented on code in PR #28714: URL: https://github.com/apache/flink/pull/28714#discussion_r3572792668
########## flink-tests/src/test/java/org/apache/flink/test/checkpointing/ChainingMaxParallelismStateLossITCase.java: ########## @@ -0,0 +1,346 @@ +/* + * 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 java.util.concurrent.atomic.AtomicInteger; + +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.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Reproduces keyed-state corruption caused by operator chaining silently overriding a keyed + * operator's explicit max parallelism across a savepoint/restore. + * + * <p>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) silently replaces the operator's + * explicit one. The savepoint's key-group count therefore differs from the restore vertex's, and + * the direction decides the symptom: + * + * <ul> + * <li>explicit < head-derived: the saved key groups fit the restored range, so restore + * succeeds but keys route to different key groups than they were stored under -> per-key + * state is silently lost. + * <li>explicit > head-derived: the saved key groups fall outside the restored range, so + * restore fails with {@code IllegalStateException: The key group must belong to the backend}. + * </ul> + */ +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<Integer, Long> COUNTS = new ConcurrentHashMap<>(); + + /** + * The max parallelism the keyed operator actually ran with, published from its {@code open()}. + */ + private static final AtomicInteger RESTORE_EFFECTIVE_MAX_PARALLELISM = new AtomicInteger(-1); + + 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 silentlyLosesKeyedStateWhenExplicitMaxParallelismBelowChainHead(String backend) + throws Exception { + startCluster(backend); + + final Map<Integer, Long> finalCounts = + savepointChainedOffRestoreChainedOn(EXPLICIT_BELOW_HEAD); + + // The keyed operator ran at the chain head's max parallelism, not its own explicit value. + assertThat(RESTORE_EFFECTIVE_MAX_PARALLELISM.get()).isEqualTo(CHAIN_HEAD_MAX_PARALLELISM); + // Every key's state was lost: counts restart from zero (job 2 only) rather than continuing. + assertThat(finalCounts).hasSize(NUM_KEYS); + assertThat(finalCounts.values()).allMatch(count -> count == JOB2_PER_KEY); + } + + @ParameterizedTest(name = "backend={0}") + @ValueSource(strings = {"hashmap", "rocksdb"}) + void failsRestoreWhenExplicitMaxParallelismAboveChainHead(String backend) throws Exception { + startCluster(backend); + + assertThatThrownBy(() -> savepointChainedOffRestoreChainedOn(EXPLICIT_ABOVE_HEAD)) + .hasStackTraceContaining("The key group must belong to the backend"); + } + + /** + * Runs one savepoint (chaining OFF, keyed operator keeps its explicit max parallelism) then + * restore (chaining ON, the explicit value is overridden by the chain head's) cycle, and + * returns the per-key counts after the restore job. + */ + private Map<Integer, Long> savepointChainedOffRestoreChainedOn(int keyedMaxParallelism) + throws Exception { + COUNTS.clear(); + RESTORE_EFFECTIVE_MAX_PARALLELISM.set(-1); + + 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): restore and feed JOB2_PER_KEY more per key. + COUNTS.clear(); + RESTORE_EFFECTIVE_MAX_PARALLELISM.set(-1); + 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<Integer> 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<Integer, Integer> keyed = + DataStreamUtils.reinterpretAsKeyedStream( + source, (KeySelector<Integer, Integer>) value -> value % NUM_KEYS); Review Comment: Yes, I agree and don't typically don't prefer using it. I went down that road and tried iterating several times with just a plain `keyBy` to trigger the issue with very no luck. Per Claude: > keyBy inserts a hash (KeyGroupStreamPartitioner) edge, which breaks the chain, so the keyed operator always ends up as a chain head with its own maxParallelism. There's no second operator on the vertex to disagree with, and on restore it just adopts its own recorded value. To hit this the keyed operator has to be a chained non-head sharing a vertex with something that recorded a different maxParallelism, and reinterpretAsKeyedStream (forward edge) is the only pure-streaming way to get that topology. I'll try exploring a few other ways to reproduce this issue without it and report back. 👍 ########## flink-tests/src/test/java/org/apache/flink/test/checkpointing/ChainingMaxParallelismStateLossITCase.java: ########## @@ -0,0 +1,346 @@ +/* + * 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 java.util.concurrent.atomic.AtomicInteger; + +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.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Reproduces keyed-state corruption caused by operator chaining silently overriding a keyed + * operator's explicit max parallelism across a savepoint/restore. + * + * <p>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) silently replaces the operator's + * explicit one. The savepoint's key-group count therefore differs from the restore vertex's, and + * the direction decides the symptom: + * + * <ul> + * <li>explicit < head-derived: the saved key groups fit the restored range, so restore + * succeeds but keys route to different key groups than they were stored under -> per-key + * state is silently lost. + * <li>explicit > head-derived: the saved key groups fall outside the restored range, so + * restore fails with {@code IllegalStateException: The key group must belong to the backend}. + * </ul> + */ +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<Integer, Long> COUNTS = new ConcurrentHashMap<>(); + + /** + * The max parallelism the keyed operator actually ran with, published from its {@code open()}. + */ + private static final AtomicInteger RESTORE_EFFECTIVE_MAX_PARALLELISM = new AtomicInteger(-1); + + 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 silentlyLosesKeyedStateWhenExplicitMaxParallelismBelowChainHead(String backend) + throws Exception { + startCluster(backend); + + final Map<Integer, Long> finalCounts = + savepointChainedOffRestoreChainedOn(EXPLICIT_BELOW_HEAD); + + // The keyed operator ran at the chain head's max parallelism, not its own explicit value. + assertThat(RESTORE_EFFECTIVE_MAX_PARALLELISM.get()).isEqualTo(CHAIN_HEAD_MAX_PARALLELISM); + // Every key's state was lost: counts restart from zero (job 2 only) rather than continuing. + assertThat(finalCounts).hasSize(NUM_KEYS); + assertThat(finalCounts.values()).allMatch(count -> count == JOB2_PER_KEY); + } + + @ParameterizedTest(name = "backend={0}") + @ValueSource(strings = {"hashmap", "rocksdb"}) + void failsRestoreWhenExplicitMaxParallelismAboveChainHead(String backend) throws Exception { + startCluster(backend); + + assertThatThrownBy(() -> savepointChainedOffRestoreChainedOn(EXPLICIT_ABOVE_HEAD)) + .hasStackTraceContaining("The key group must belong to the backend"); + } + + /** + * Runs one savepoint (chaining OFF, keyed operator keeps its explicit max parallelism) then + * restore (chaining ON, the explicit value is overridden by the chain head's) cycle, and + * returns the per-key counts after the restore job. + */ + private Map<Integer, Long> savepointChainedOffRestoreChainedOn(int keyedMaxParallelism) + throws Exception { + COUNTS.clear(); + RESTORE_EFFECTIVE_MAX_PARALLELISM.set(-1); + + 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): restore and feed JOB2_PER_KEY more per key. + COUNTS.clear(); + RESTORE_EFFECTIVE_MAX_PARALLELISM.set(-1); + 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<Integer> 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<Integer, Integer> keyed = + DataStreamUtils.reinterpretAsKeyedStream( + source, (KeySelector<Integer, Integer>) value -> value % NUM_KEYS); + + final SingleOutputStreamOperator<Tuple2<Integer, Long>> 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(); + } Review Comment: Yup, this is fair. I can move it into the job configuration instead. 👍 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
