From 9ed1327a767b4b9e6dcef0a7f820b9e8eba2fbbe Mon Sep 17 00:00:00 2001 From: Rui Fan <1996fanrui@gmail.com> Date: Fri, 31 Jul 2026 12:34:32 +0200 Subject: [PATCH 1/3] [FLINK-40269][tests] Add same-upstream unaligned rescale ITCase --- .../runtime/testutils/CommonTestUtils.java | 13 + ...edCheckpointRescaleSameUpstreamITCase.java | 258 ++++++++++++++++++ 2 files changed, 271 insertions(+) create mode 100644 flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleSameUpstreamITCase.java 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); + } +} From 039cb7ad106e9912bffcb9d530ff498dc51552e1 Mon Sep 17 00:00:00 2001 From: Rui Fan <1996fanrui@gmail.com> Date: Fri, 31 Jul 2026 14:47:05 +0200 Subject: [PATCH 2/3] [FLINK-40269][runtime] Fix channel state assignment for duplicate connections --- .../checkpoint/TaskStateAssignment.java | 88 ++++++++++++------- .../StateAssignmentOperationTest.java | 66 ++++++++++++++ 2 files changed, 121 insertions(+), 33 deletions(-) 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..de13e4c15fd94e 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,10 +266,13 @@ public OperatorSubtaskState getSubtaskState(OperatorInstanceID instanceID) { instanceID, inputOperatorID, getUpstreamAssignments(), - (assignment, recompute) -> { + (gateIndex, assignment, recompute) -> { int assignmentIndex = - getAssignmentIndex( - assignment.getDownstreamAssignments(), this); + assignment.findResultPartitionIndex( + executionJobVertex + .getInputs() + .get(gateIndex) + .getId()); return assignment.getOutputMapping(assignmentIndex, recompute); }, inputSubtaskMappings, @@ -320,9 +317,11 @@ private InflightDataRescalingDescriptor computeOutputRescalingDescriptor( instanceID, outputOperatorID, getDownstreamAssignments(), - (downstreamAssignment, recompute) -> { + (partitionIndex, downstreamAssignment, recompute) -> { int assignmentIndex = - getAssignmentIndex(downstreamAssignment.getUpstreamAssignments(), this); + downstreamAssignment.findInputGateIndex( + executionJobVertex.getProducedDataSets()[partitionIndex] + .getId()); return downstreamAssignment.getInputMapping(assignmentIndex, recompute); }, outputSubtaskMappings, @@ -355,7 +354,7 @@ private InflightDataRescalingDescriptor createRescalingDescriptor( OperatorInstanceID instanceID, OperatorID expectedOperatorID, TaskStateAssignment[] connectedAssignments, - BiFunction mappingRetriever, + MappingRetriever mappingRetriever, Map subtaskGateOrPartitionMappings, Function subtaskMappingCalculator, boolean isInput) { @@ -364,8 +363,11 @@ private InflightDataRescalingDescriptor createRescalingDescriptor( } SubtasksRescaleMapping[] rescaledChannelsMappings = - Arrays.stream(connectedAssignments) - .map(assignment -> mappingRetriever.apply(assignment, false)) + IntStream.range(0, connectedAssignments.length) + .mapToObj( + index -> + mappingRetriever.get( + index, connectedAssignments[index], false)) .toArray(SubtasksRescaleMapping[]::new); // no state on input and output, especially for any aligned checkpoint @@ -378,7 +380,7 @@ private InflightDataRescalingDescriptor createRescalingDescriptor( createGateOrPartitionRescalingDescriptors( instanceID, connectedAssignments, - assignment -> mappingRetriever.apply(assignment, true), + mappingRetriever, subtaskGateOrPartitionMappings, subtaskMappingCalculator, rescaledChannelsMappings, @@ -398,7 +400,7 @@ private InflightDataRescalingDescriptor createRescalingDescriptor( createGateOrPartitionRescalingDescriptors( OperatorInstanceID instanceID, TaskStateAssignment[] connectedAssignments, - Function mappingCalculator, + MappingRetriever mappingRetriever, Map subtaskGateOrPartitionMappings, Function subtaskMappingCalculator, SubtasksRescaleMapping[] rescaledChannelsMappings, @@ -415,8 +417,10 @@ private InflightDataRescalingDescriptor createRescalingDescriptor( Optional.ofNullable(rescaledChannelsMappings[partition]) .orElseGet( () -> - mappingCalculator.apply( - connectedAssignment)); + mappingRetriever.get( + partition, + connectedAssignment, + true)); SubtasksRescaleMapping subtaskMapping = Optional.ofNullable( subtaskGateOrPartitionMappings.get(partition)) @@ -547,12 +551,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 +571,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 +638,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 @@ -702,4 +718,10 @@ public boolean isMayHaveAmbiguousSubtasks() { return mayHaveAmbiguousSubtasks; } } + + @FunctionalInterface + private interface MappingRetriever { + SubtasksRescaleMapping get( + int connectedIndex, TaskStateAssignment connectedAssignment, boolean recompute); + } } 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, From 26be3b10bc8b403ad0994a743e6b9ac27b1fe320 Mon Sep 17 00:00:00 2001 From: Roman Khachatryan Date: Fri, 31 Jul 2026 19:47:17 +0200 Subject: [PATCH 3/3] [hotfix][runtime] Minor simplification of TaskStateAssignment Resolve connected mappings via isInput instead of a MappingRetriever lambda --- .../checkpoint/TaskStateAssignment.java | 59 ++++++++++--------- 1 file changed, 31 insertions(+), 28 deletions(-) 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 de13e4c15fd94e..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 @@ -266,15 +266,6 @@ public OperatorSubtaskState getSubtaskState(OperatorInstanceID instanceID) { instanceID, inputOperatorID, getUpstreamAssignments(), - (gateIndex, assignment, recompute) -> { - int assignmentIndex = - assignment.findResultPartitionIndex( - executionJobVertex - .getInputs() - .get(gateIndex) - .getId()); - return assignment.getOutputMapping(assignmentIndex, recompute); - }, inputSubtaskMappings, this::getInputMapping, true)) @@ -317,13 +308,6 @@ private InflightDataRescalingDescriptor computeOutputRescalingDescriptor( instanceID, outputOperatorID, getDownstreamAssignments(), - (partitionIndex, downstreamAssignment, recompute) -> { - int assignmentIndex = - downstreamAssignment.findInputGateIndex( - executionJobVertex.getProducedDataSets()[partitionIndex] - .getId()); - return downstreamAssignment.getInputMapping(assignmentIndex, recompute); - }, outputSubtaskMappings, this::getOutputMapping, false); @@ -354,7 +338,6 @@ private InflightDataRescalingDescriptor createRescalingDescriptor( OperatorInstanceID instanceID, OperatorID expectedOperatorID, TaskStateAssignment[] connectedAssignments, - MappingRetriever mappingRetriever, Map subtaskGateOrPartitionMappings, Function subtaskMappingCalculator, boolean isInput) { @@ -366,8 +349,8 @@ private InflightDataRescalingDescriptor createRescalingDescriptor( IntStream.range(0, connectedAssignments.length) .mapToObj( index -> - mappingRetriever.get( - index, connectedAssignments[index], false)) + getConnectedMapping( + isInput, index, connectedAssignments[index], false)) .toArray(SubtasksRescaleMapping[]::new); // no state on input and output, especially for any aligned checkpoint @@ -380,7 +363,6 @@ private InflightDataRescalingDescriptor createRescalingDescriptor( createGateOrPartitionRescalingDescriptors( instanceID, connectedAssignments, - mappingRetriever, subtaskGateOrPartitionMappings, subtaskMappingCalculator, rescaledChannelsMappings, @@ -400,7 +382,6 @@ private InflightDataRescalingDescriptor createRescalingDescriptor( createGateOrPartitionRescalingDescriptors( OperatorInstanceID instanceID, TaskStateAssignment[] connectedAssignments, - MappingRetriever mappingRetriever, Map subtaskGateOrPartitionMappings, Function subtaskMappingCalculator, SubtasksRescaleMapping[] rescaledChannelsMappings, @@ -417,7 +398,8 @@ private InflightDataRescalingDescriptor createRescalingDescriptor( Optional.ofNullable(rescaledChannelsMappings[partition]) .orElseGet( () -> - mappingRetriever.get( + getConnectedMapping( + isInput, partition, connectedAssignment, true)); @@ -489,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) { @@ -498,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]; @@ -718,10 +727,4 @@ public boolean isMayHaveAmbiguousSubtasks() { return mayHaveAmbiguousSubtasks; } } - - @FunctionalInterface - private interface MappingRetriever { - SubtasksRescaleMapping get( - int connectedIndex, TaskStateAssignment connectedAssignment, boolean recompute); - } }