pnowojski commented on code in PR #28714: URL: https://github.com/apache/flink/pull/28714#discussion_r3571901153
########## 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: Can this issue be reproduced without `reinterpretAsKeyedStream`? 🤔 Just normal `keyBy` followed by a chained/not chained operator? `reinterpretAsKeyedStream` is a bit hacky and doesn't represents the normal use cases. It's intention is if partitions in your sources are already keyed by something, so you don't need to `keyBy` that data stream again. Here it only works because you constrain the parallelism to `1` which is also a bit artificial constraint. ########## 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: Why do you need this manually started instead of using ` @ClassRule` on the `MiniClusterWithClientResource` field? There are plenty of examples when such cluster has a custom configuration. ########## flink-tests/src/test/java/org/apache/flink/test/checkpointing/ChainingMaxParallelismStateLossITCase.java: ########## @@ -58,32 +58,23 @@ 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. + * Verifies that restoring a savepoint is rejected when a chaining change places operators with + * different recorded max parallelism onto a single keyed vertex. Review Comment: Some commits needs to be squashed ########## 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)); Review Comment: Why do we need a deadline? Our default convetion is to just wait indefinetly: - often even reasonable deadlines are violated due to slow CI - if something indeed deadlocks in CI, we have a global timeout of ~4h, after which CI captures thread dump. By adding custom deadlines/timeouts, you are basically bypassing that thread dump capture mechanism, which often is the only way one can analyse a deadlock (especially painful if it's extremely rare and can't be reproduced locally) -- 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]
