diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskStateAssignment.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskStateAssignment.java index a6db5e4837f859..dc663f01bd0faa 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskStateAssignment.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskStateAssignment.java @@ -53,7 +53,6 @@ import java.util.Objects; import java.util.Optional; import java.util.Set; -import java.util.function.BiFunction; import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -230,11 +229,6 @@ public TaskStateAssignment[] getDownstreamAssignments() { return downstreamAssignments; } - private static int getAssignmentIndex( - TaskStateAssignment[] assignments, TaskStateAssignment assignment) { - return Arrays.asList(assignments).indexOf(assignment); - } - public TaskStateAssignment[] getUpstreamAssignments() { if (upstreamAssignments == null) { upstreamAssignments = @@ -272,12 +266,6 @@ public OperatorSubtaskState getSubtaskState(OperatorInstanceID instanceID) { instanceID, inputOperatorID, getUpstreamAssignments(), - (assignment, recompute) -> { - int assignmentIndex = - getAssignmentIndex( - assignment.getDownstreamAssignments(), this); - return assignment.getOutputMapping(assignmentIndex, recompute); - }, inputSubtaskMappings, this::getInputMapping, true)) @@ -320,11 +308,6 @@ private InflightDataRescalingDescriptor computeOutputRescalingDescriptor( instanceID, outputOperatorID, getDownstreamAssignments(), - (downstreamAssignment, recompute) -> { - int assignmentIndex = - getAssignmentIndex(downstreamAssignment.getUpstreamAssignments(), this); - return downstreamAssignment.getInputMapping(assignmentIndex, recompute); - }, outputSubtaskMappings, this::getOutputMapping, false); @@ -355,7 +338,6 @@ private InflightDataRescalingDescriptor createRescalingDescriptor( OperatorInstanceID instanceID, OperatorID expectedOperatorID, TaskStateAssignment[] connectedAssignments, - BiFunction mappingRetriever, Map subtaskGateOrPartitionMappings, Function subtaskMappingCalculator, boolean isInput) { @@ -364,8 +346,11 @@ private InflightDataRescalingDescriptor createRescalingDescriptor( } SubtasksRescaleMapping[] rescaledChannelsMappings = - Arrays.stream(connectedAssignments) - .map(assignment -> mappingRetriever.apply(assignment, false)) + IntStream.range(0, connectedAssignments.length) + .mapToObj( + index -> + getConnectedMapping( + isInput, index, connectedAssignments[index], false)) .toArray(SubtasksRescaleMapping[]::new); // no state on input and output, especially for any aligned checkpoint @@ -378,7 +363,6 @@ private InflightDataRescalingDescriptor createRescalingDescriptor( createGateOrPartitionRescalingDescriptors( instanceID, connectedAssignments, - assignment -> mappingRetriever.apply(assignment, true), subtaskGateOrPartitionMappings, subtaskMappingCalculator, rescaledChannelsMappings, @@ -398,7 +382,6 @@ private InflightDataRescalingDescriptor createRescalingDescriptor( createGateOrPartitionRescalingDescriptors( OperatorInstanceID instanceID, TaskStateAssignment[] connectedAssignments, - Function mappingCalculator, Map subtaskGateOrPartitionMappings, Function subtaskMappingCalculator, SubtasksRescaleMapping[] rescaledChannelsMappings, @@ -415,8 +398,11 @@ private InflightDataRescalingDescriptor createRescalingDescriptor( Optional.ofNullable(rescaledChannelsMappings[partition]) .orElseGet( () -> - mappingCalculator.apply( - connectedAssignment)); + getConnectedMapping( + isInput, + partition, + connectedAssignment, + true)); SubtasksRescaleMapping subtaskMapping = Optional.ofNullable( subtaskGateOrPartitionMappings.get(partition)) @@ -485,6 +471,11 @@ private SubtasksRescaleMapping getOutputMapping(int assignmentIndex, boolean rec } } + private SubtasksRescaleMapping getOutputMapping( + IntermediateDataSetID resultId, boolean recompute) { + return getOutputMapping(findResultPartitionIndex(resultId), recompute); + } + private SubtasksRescaleMapping getInputMapping(int assignmentIndex, boolean recompute) { SubtasksRescaleMapping mapping = inputSubtaskMappings.get(assignmentIndex); if (recompute && mapping == null) { @@ -494,6 +485,28 @@ private SubtasksRescaleMapping getInputMapping(int assignmentIndex, boolean reco } } + private SubtasksRescaleMapping getInputMapping( + IntermediateDataSetID resultId, boolean recompute) { + return getInputMapping(findInputGateIndex(resultId), recompute); + } + + /** + * Resolves the mapping on {@code connectedAssignment} that corresponds to {@code index} on + * {@code this} assignment, disambiguating by {@link IntermediateDataSetID} rather than by + * array position (multiple edges can connect the same pair of job vertices). + */ + private SubtasksRescaleMapping getConnectedMapping( + boolean isInput, int index, TaskStateAssignment connectedAssignment, boolean recompute) { + if (isInput) { + IntermediateDataSetID resultId = executionJobVertex.getInputs().get(index).getId(); + return connectedAssignment.getOutputMapping(resultId, recompute); + } else { + IntermediateDataSetID resultId = + executionJobVertex.getProducedDataSets()[index].getId(); + return connectedAssignment.getInputMapping(resultId, recompute); + } + } + public SubtasksRescaleMapping getOutputMapping(int partitionIndex) { final TaskStateAssignment downstreamAssignment = getDownstreamAssignments()[partitionIndex]; final IntermediateResult output = executionJobVertex.getProducedDataSets()[partitionIndex]; @@ -547,12 +560,8 @@ public boolean hasInFlightDataForInputGate(int gateIndex) { if (upstreamAssignment != null && upstreamAssignment.hasOutputState()) { IntermediateResult inputResult = executionJobVertex.getInputs().get(gateIndex); IntermediateDataSetID resultId = inputResult.getId(); - IntermediateResult[] producedDataSets = inputResult.getProducer().getProducedDataSets(); - for (int i = 0; i < producedDataSets.length; i++) { - if (producedDataSets[i].getId().equals(resultId)) { - return upstreamAssignment.outputStatePartitions.contains(i); - } - } + return upstreamAssignment.outputStatePartitions.contains( + upstreamAssignment.findResultPartitionIndex(resultId)); } return false; @@ -571,12 +580,8 @@ public boolean hasInFlightDataForResultPartition(int partitionIndex) { IntermediateResult producedResult = executionJobVertex.getProducedDataSets()[partitionIndex]; IntermediateDataSetID resultId = producedResult.getId(); - List inputs = downstreamAssignment.executionJobVertex.getInputs(); - for (int i = 0; i < inputs.size(); i++) { - if (inputs.get(i).getId().equals(resultId)) { - return downstreamAssignment.inputStateGates.contains(i); - } - } + return downstreamAssignment.inputStateGates.contains( + downstreamAssignment.findInputGateIndex(resultId)); } return false; } @@ -642,15 +647,35 @@ private int findInputGateIdxForResultPartition(int partitionIndex) { IntermediateResult producedResult = executionJobVertex.getProducedDataSets()[partitionIndex]; - IntermediateDataSetID resultId = producedResult.getId(); - List inputs = downstreamAssignment.executionJobVertex.getInputs(); + return downstreamAssignment.findInputGateIndex(producedResult.getId()); + } + + private int findInputGateIndex(IntermediateDataSetID resultId) { + List inputs = executionJobVertex.getInputs(); for (int i = 0; i < inputs.size(); i++) { if (inputs.get(i).getId().equals(resultId)) { return i; } } throw new IllegalArgumentException( - "No channel rescaler found during rescaling of channel state"); + "No input gate found for intermediate data set " + + resultId + + " in " + + executionJobVertex.getName()); + } + + private int findResultPartitionIndex(IntermediateDataSetID resultId) { + IntermediateResult[] producedDataSets = executionJobVertex.getProducedDataSets(); + for (int i = 0; i < producedDataSets.length; i++) { + if (producedDataSets[i].getId().equals(resultId)) { + return i; + } + } + throw new IllegalArgumentException( + "No result partition found for intermediate data set " + + resultId + + " in " + + executionJobVertex.getName()); } @Override 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 1f2c9c7a3031c1..921ffb63060831 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 @@ -581,6 +581,72 @@ void testChannelStateAssignmentDownscalingTwoDifferentGates() RESCALING)))); } + @Test + void testChannelStateAssignmentUsesResultIdForDuplicateJobVertexConnections() + throws JobException, JobExecutionException { + int oldParallelism = 3; + int newParallelism = 2; + JobVertex upstream = createJobVertex(new OperatorID(), newParallelism); + JobVertex downstream = createJobVertex(new OperatorID(), newParallelism); + OperatorID upstreamOperator = upstream.getOperatorIDs().get(0).getGeneratedOperatorID(); + OperatorID downstreamOperator = downstream.getOperatorIDs().get(0).getGeneratedOperatorID(); + Random random = new Random(); + + OperatorState upstreamState = + new OperatorState("", "", upstreamOperator, oldParallelism, MAX_P); + OperatorState downstreamState = + new OperatorState("", "", downstreamOperator, oldParallelism, MAX_P); + for (int i = 0; i < oldParallelism; i++) { + upstreamState.putState( + i, + OperatorSubtaskState.builder() + .setResultSubpartitionState( + new StateObjectCollection<>( + asList( + createNewResultSubpartitionStateHandle( + 10, 0, random), + createNewResultSubpartitionStateHandle( + 10, 1, random)))) + .build()); + downstreamState.putState( + i, + OperatorSubtaskState.builder() + .setInputChannelState( + new StateObjectCollection<>( + asList( + createNewInputChannelStateHandle(10, 0, random), + createNewInputChannelStateHandle( + 10, 1, random)))) + .build()); + } + Map states = new HashMap<>(); + states.put(upstreamOperator, upstreamState); + states.put(downstreamOperator, downstreamState); + + connectVertices(upstream, downstream, RANGE, RANGE); + connectVertices(upstream, downstream, ROUND_ROBIN, ROUND_ROBIN); + + Map vertices = toExecutionVertices(upstream, downstream); + + new StateAssignmentOperation(0, new HashSet<>(vertices.values()), states, false, false) + .assignStates(); + + InflightDataRescalingDescriptor outputDescriptor = + getAssignedState(vertices.get(upstreamOperator), upstreamOperator, 0) + .getOutputRescalingDescriptor(); + InflightDataRescalingDescriptor inputDescriptor = + getAssignedState(vertices.get(downstreamOperator), downstreamOperator, 0) + .getInputRescalingDescriptor(); + assertThat(outputDescriptor.getChannelMapping(0)) + .isEqualTo(RANGE.getNewToOldSubtasksMapping(oldParallelism, newParallelism)); + assertThat(outputDescriptor.getChannelMapping(1)) + .isEqualTo(ROUND_ROBIN.getNewToOldSubtasksMapping(oldParallelism, newParallelism)); + assertThat(inputDescriptor.getChannelMapping(0)) + .isEqualTo(RANGE.getNewToOldSubtasksMapping(oldParallelism, newParallelism)); + assertThat(inputDescriptor.getChannelMapping(1)) + .isEqualTo(ROUND_ROBIN.getNewToOldSubtasksMapping(oldParallelism, newParallelism)); + } + private InflightDataGateOrPartitionRescalingDescriptor gate( int[] oldIndices, RescaleMappings rescaleMapping, diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/testutils/CommonTestUtils.java b/flink-runtime/src/test/java/org/apache/flink/runtime/testutils/CommonTestUtils.java index e70ac372bcd86e..557c49395ca1db 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/testutils/CommonTestUtils.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/testutils/CommonTestUtils.java @@ -345,6 +345,15 @@ public static void waitForSubtasksToFinish( /** Wait for one checkpoint with in-flight buffers. */ public static String waitForCheckpointWithInflightBuffers(JobID jobID, MiniCluster miniCluster) throws Exception { + return waitForCheckpointWithInflightBuffers(jobID, miniCluster, 1); + } + + /** + * Wait for at least {@code minCompletedCheckpoints} completed checkpoints and return the latest + * checkpoint with in-flight buffers. + */ + public static String waitForCheckpointWithInflightBuffers( + JobID jobID, MiniCluster miniCluster, long minCompletedCheckpoints) throws Exception { CompletableFuture checkpointPath = new CompletableFuture<>(); waitForCheckpoints( jobID, @@ -353,6 +362,10 @@ public static String waitForCheckpointWithInflightBuffers(JobID jobID, MiniClust if (checkpointStatsSnapshot == null) { return false; } + if (checkpointStatsSnapshot.getCounts().getNumberOfCompletedCheckpoints() + < minCompletedCheckpoints) { + return false; + } CompletedCheckpointStats latestCompletedCheckpoint = checkpointStatsSnapshot.getHistory().getLatestCompletedCheckpoint(); diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleSameUpstreamITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleSameUpstreamITCase.java new file mode 100644 index 00000000000000..15ef9366ff3a3a --- /dev/null +++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleSameUpstreamITCase.java @@ -0,0 +1,258 @@ +/* + * 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.JobStatus; +import org.apache.flink.api.common.functions.MapFunction; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.configuration.CheckpointingOptions; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.ExternalizedCheckpointRetention; +import org.apache.flink.configuration.MemorySize; +import org.apache.flink.configuration.RestartStrategyOptions; +import org.apache.flink.configuration.StateRecoveryOptions; +import org.apache.flink.configuration.TaskManagerOptions; +import org.apache.flink.core.execution.JobClient; +import org.apache.flink.runtime.jobgraph.JobGraph; +import org.apache.flink.runtime.minicluster.MiniCluster; +import org.apache.flink.runtime.minicluster.MiniClusterJobClient; +import org.apache.flink.runtime.testutils.CommonTestUtils; +import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration; +import org.apache.flink.streaming.api.datastream.DataStream; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.api.functions.co.CoMapFunction; +import org.apache.flink.streaming.api.functions.sink.v2.DiscardingSink; +import org.apache.flink.test.junit5.InjectMiniCluster; +import org.apache.flink.test.junit5.MiniClusterExtension; +import org.apache.flink.testutils.junit.extensions.parameterized.Parameter; +import org.apache.flink.testutils.junit.extensions.parameterized.ParameterizedTestExtension; +import org.apache.flink.testutils.junit.extensions.parameterized.Parameters; + +import org.junit.jupiter.api.TestTemplate; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.api.io.TempDir; + +import javax.annotation.Nullable; + +import java.io.File; +import java.time.Duration; +import java.util.Collections; + +import static org.apache.flink.configuration.RestartStrategyOptions.RestartStrategyType.NO_RESTART_STRATEGY; + +/** + * Integration test for unaligned checkpoint rescaling when one downstream task receives multiple + * inputs derived from the same upstream stream. + */ +@ExtendWith(ParameterizedTestExtension.class) +class UnalignedCheckpointRescaleSameUpstreamITCase { + + private static final int INITIAL_PARALLELISM = 2; + private static final int RESTORED_PARALLELISM = 4; + private static final int SLOTS_PER_TASK_MANAGER = 8; + private static final int CHECKPOINTS_TO_WAIT = 10; + + @RegisterExtension + private static final MiniClusterExtension MINI_CLUSTER_EXTENSION = + new MiniClusterExtension( + new MiniClusterResourceConfiguration.Builder() + .setConfiguration( + new Configuration() + .set(CheckpointingOptions.MAX_RETAINED_CHECKPOINTS, 50)) + .setNumberTaskManagers(1) + .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER) + .build()); + + @TempDir private File temporaryFolder; + + @Parameter private boolean recoverOutputOnDownstream; + + @Parameter(1) + private SameUpstreamDag dag; + + @Parameters(name = "recoverOutputOnDownstream={0}, dag={1}") + private static Object[][] parameters() { + return new Object[][] { + new Object[] {false, SameUpstreamDag.CONNECT}, + new Object[] {true, SameUpstreamDag.CONNECT}, + new Object[] {false, SameUpstreamDag.UNION}, + new Object[] {true, SameUpstreamDag.UNION} + }; + } + + @TestTemplate + void testRescaleFromUnalignedCheckpointWithSameUpstream( + @InjectMiniCluster MiniCluster miniCluster) throws Exception { + final JobGraph initialJobGraph = + createJobGraph(null, INITIAL_PARALLELISM, recoverOutputOnDownstream); + + final JobClient initialJobClient = submitJob(initialJobGraph, miniCluster); + final String checkpointPath; + try { + waitForRunning(initialJobClient, miniCluster); + checkpointPath = + CommonTestUtils.waitForCheckpointWithInflightBuffers( + initialJobGraph.getJobID(), miniCluster, CHECKPOINTS_TO_WAIT); + } finally { + initialJobClient.cancel().get(); + } + + final JobGraph restoredJobGraph = + createJobGraph(checkpointPath, RESTORED_PARALLELISM, recoverOutputOnDownstream); + + final JobClient restoredJobClient = submitJob(restoredJobGraph, miniCluster); + try { + waitForRunning(restoredJobClient, miniCluster); + CommonTestUtils.waitForCheckpointWithInflightBuffers( + restoredJobGraph.getJobID(), miniCluster); + } finally { + cancelIfRunning(restoredJobClient); + } + } + + private JobGraph createJobGraph( + @Nullable String recoveryPath, int parallelism, boolean recoverOutputOnDownstream) + throws Exception { + final StreamExecutionEnvironment env = + StreamExecutionEnvironment.getExecutionEnvironment( + createConfiguration(recoveryPath, recoverOutputOnDownstream)); + env.disableOperatorChaining(); + dag.build(env, parallelism); + return env.getStreamGraph().getJobGraph(); + } + + private Configuration createConfiguration( + @Nullable String recoveryPath, boolean recoverOutputOnDownstream) { + final Configuration conf = new Configuration(); + conf.set(CheckpointingOptions.CHECKPOINTING_INTERVAL, Duration.ofSeconds(1)); + conf.set(CheckpointingOptions.ALIGNED_CHECKPOINT_TIMEOUT, Duration.ofSeconds(0)); + conf.set(RestartStrategyOptions.RESTART_STRATEGY, NO_RESTART_STRATEGY.getMainValue()); + conf.set( + CheckpointingOptions.EXTERNALIZED_CHECKPOINT_RETENTION, + ExternalizedCheckpointRetention.RETAIN_ON_CANCELLATION); + conf.set(CheckpointingOptions.CHECKPOINTS_DIRECTORY, temporaryFolder.toURI().toString()); + conf.set(CheckpointingOptions.ENABLE_UNALIGNED, true); + conf.set( + CheckpointingOptions.UNALIGNED_RECOVER_OUTPUT_ON_DOWNSTREAM, + recoverOutputOnDownstream); + conf.set(TaskManagerOptions.MEMORY_SEGMENT_SIZE, MemorySize.parse("4 kb")); + if (recoveryPath != null) { + conf.set(StateRecoveryOptions.SAVEPOINT_PATH, recoveryPath); + } + return conf; + } + + private static JobClient submitJob(JobGraph jobGraph, MiniCluster miniCluster) + throws Exception { + miniCluster.submitJob(jobGraph).get(); + return new MiniClusterJobClient( + jobGraph.getJobID(), + miniCluster, + Thread.currentThread().getContextClassLoader(), + MiniClusterJobClient.JobFinalizationBehavior.NOTHING); + } + + private static void waitForRunning(JobClient jobClient, MiniCluster miniCluster) + throws Exception { + CommonTestUtils.waitForJobStatus(jobClient, Collections.singletonList(JobStatus.RUNNING)); + CommonTestUtils.waitForAllTaskRunning(miniCluster, jobClient.getJobID(), false); + } + + private static void cancelIfRunning(JobClient jobClient) throws Exception { + if (jobClient.getJobStatus().get() != JobStatus.FAILED) { + jobClient.cancel().get(); + } + } + + private static class SleepingCoMap implements CoMapFunction { + @Override + public T map1(T value) throws Exception { + Thread.sleep(1); + return value; + } + + @Override + public T map2(T value) throws Exception { + Thread.sleep(5); + return value; + } + } + + private static class SleepingMap implements MapFunction { + @Override + public T map(T value) throws Exception { + Thread.sleep(5); + return value; + } + } + + private enum SameUpstreamDag { + CONNECT { + @Override + void build(StreamExecutionEnvironment env, int parallelism) { + final DataStream upstream = + env.fromSequence(0, Long.MAX_VALUE) + .name("Upstream") + .uid("upstream") + .setParallelism(parallelism); + final DataStream leftInput = upstream.rebalance(); + final DataStream rightInput = + upstream.keyBy((KeySelector) value -> value); + + leftInput + .connect(rightInput) + .map(new SleepingCoMap<>()) + .name("Co-Map") + .uid("co-map") + .setParallelism(parallelism) + .sinkTo(new DiscardingSink<>()) + .name("Discarding Sink") + .uid("sink") + .setParallelism(parallelism); + } + }, + UNION { + @Override + void build(StreamExecutionEnvironment env, int parallelism) { + final DataStream upstream = + env.fromSequence(0, Long.MAX_VALUE) + .name("Upstream") + .uid("upstream") + .setParallelism(parallelism); + final DataStream leftInput = upstream.rebalance(); + final DataStream rightInput = + upstream.keyBy((KeySelector) value -> value); + + leftInput + .union(rightInput) + .map(new SleepingMap<>()) + .name("Slow Map") + .uid("slow-map") + .setParallelism(parallelism) + .sinkTo(new DiscardingSink<>()) + .name("Discarding Sink") + .uid("sink") + .setParallelism(parallelism); + } + }; + + abstract void build(StreamExecutionEnvironment env, int parallelism); + } +}