rionmonster commented on code in PR #28714: URL: https://github.com/apache/flink/pull/28714#discussion_r3572793944
########## 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: I can go ahead and drop this as well to clean things up. ########## 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: Will do! I'll squash it down to a single primary commit. -- 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]
