You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by fa...@apache.org on 2025/08/29 10:38:17 UTC
(flink) 02/03: [FLINK-38267][checkpoint] Only call channel state rescaling logic for exchange with channel state to avoid UnsupportedOperationException
This is an automated email from the ASF dual-hosted git repository.
fanrui pushed a commit to branch release-2.1
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 4f17abd871524b70737365e155aef1abfacfe241
Author: Rui Fan <19...@gmail.com>
AuthorDate: Tue Aug 19 21:18:51 2025 +0200
[FLINK-38267][checkpoint] Only call channel state rescaling logic for exchange with channel state to avoid UnsupportedOperationException
---
.../InflightDataRescalingDescriptor.java | 35 +-
.../checkpoint/StateAssignmentOperation.java | 6 +
.../runtime/checkpoint/TaskStateAssignment.java | 71 +++-
.../InflightDataRescalingDescriptorTest.java | 121 +++++++
.../checkpoint/StateAssignmentOperationTest.java | 298 ++++++++++++++++-
.../runtime/checkpoint/StateHandleDummyUtil.java | 14 +-
.../flink/runtime/testutils/CommonTestUtils.java | 25 ++
...dCheckpointRescaleWithMixedExchangesITCase.java | 361 +++++++++++++++++++++
8 files changed, 921 insertions(+), 10 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/InflightDataRescalingDescriptor.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/InflightDataRescalingDescriptor.java
index 53d479250ac..4dffeee0bb9 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/InflightDataRescalingDescriptor.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/InflightDataRescalingDescriptor.java
@@ -20,6 +20,7 @@ package org.apache.flink.runtime.checkpoint;
import java.io.ObjectStreamException;
import java.io.Serializable;
import java.util.Arrays;
+import java.util.Collections;
import java.util.Objects;
import java.util.Set;
@@ -45,11 +46,11 @@ public class InflightDataRescalingDescriptor implements Serializable {
}
public int[] getOldSubtaskIndexes(int gateOrPartitionIndex) {
- return gateOrPartitionDescriptors[gateOrPartitionIndex].oldSubtaskIndexes;
+ return gateOrPartitionDescriptors[gateOrPartitionIndex].getOldSubtaskInstances();
}
public RescaleMappings getChannelMapping(int gateOrPartitionIndex) {
- return gateOrPartitionDescriptors[gateOrPartitionIndex].rescaledChannelsMappings;
+ return gateOrPartitionDescriptors[gateOrPartitionIndex].getRescaleMappings();
}
public boolean isAmbiguous(int gateOrPartitionIndex, int oldSubtaskIndex) {
@@ -112,6 +113,28 @@ public class InflightDataRescalingDescriptor implements Serializable {
*/
public static class InflightDataGateOrPartitionRescalingDescriptor implements Serializable {
+ public static final InflightDataGateOrPartitionRescalingDescriptor NO_STATE =
+ new InflightDataGateOrPartitionRescalingDescriptor(
+ new int[0],
+ RescaleMappings.identity(0, 0),
+ Collections.emptySet(),
+ MappingType.IDENTITY) {
+
+ private static final long serialVersionUID = 1L;
+
+ @Override
+ public int[] getOldSubtaskInstances() {
+ throw new UnsupportedOperationException(
+ "Cannot get old subtasks from a descriptor that represents no state.");
+ }
+
+ @Override
+ public RescaleMappings getRescaleMappings() {
+ throw new UnsupportedOperationException(
+ "Cannot get rescale mappings from a descriptor that represents no state.");
+ }
+ };
+
private static final long serialVersionUID = 1L;
/** Set when several operator instances are merged into one. */
@@ -145,6 +168,14 @@ public class InflightDataRescalingDescriptor implements Serializable {
this.mappingType = mappingType;
}
+ public int[] getOldSubtaskInstances() {
+ return oldSubtaskIndexes;
+ }
+
+ public RescaleMappings getRescaleMappings() {
+ return rescaledChannelsMappings;
+ }
+
public boolean isIdentity() {
return mappingType == MappingType.IDENTITY;
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperation.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperation.java
index 7b41b4981f0..70611f7b40e 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperation.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StateAssignmentOperation.java
@@ -386,6 +386,9 @@ public class StateAssignmentOperation {
// Parallelism of this vertex changed, distribute ResultSubpartitionStateHandle
// according to output mapping.
for (int partitionIndex = 0; partitionIndex < outputs.size(); partitionIndex++) {
+ if (!assignment.hasInFlightDataForResultPartition(partitionIndex)) {
+ continue;
+ }
final List<List<ResultSubpartitionStateHandle>> partitionState =
outputs.size() == 1
? outputOperatorState
@@ -466,6 +469,9 @@ public class StateAssignmentOperation {
// subtask 0 recovers data from old subtask 0 + 1 and subtask 1 recovers data from old
// subtask 1 + 2
for (int gateIndex = 0; gateIndex < inputs.size(); gateIndex++) {
+ if (!stateAssignment.hasInFlightDataForInputGate(gateIndex)) {
+ continue;
+ }
final RescaleMappings mapping =
stateAssignment.getInputMapping(gateIndex).getRescaleMappings();
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 54810ab7b55..4601b4716a1 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
@@ -252,7 +252,8 @@ class TaskStateAssignment {
return assignment.getOutputMapping(assignmentIndex, recompute);
},
inputSubtaskMappings,
- this::getInputMapping))
+ this::getInputMapping,
+ true))
.setOutputRescalingDescriptor(
createRescalingDescriptor(
instanceID,
@@ -265,7 +266,8 @@ class TaskStateAssignment {
return assignment.getInputMapping(assignmentIndex, recompute);
},
outputSubtaskMappings,
- this::getOutputMapping))
+ this::getOutputMapping,
+ false))
.build();
}
@@ -314,7 +316,8 @@ class TaskStateAssignment {
TaskStateAssignment[] connectedAssignments,
BiFunction<TaskStateAssignment, Boolean, SubtasksRescaleMapping> mappingRetriever,
Map<Integer, SubtasksRescaleMapping> subtaskGateOrPartitionMappings,
- Function<Integer, SubtasksRescaleMapping> subtaskMappingCalculator) {
+ Function<Integer, SubtasksRescaleMapping> subtaskMappingCalculator,
+ boolean isInput) {
if (!expectedOperatorID.equals(instanceID.getOperatorId())) {
return InflightDataRescalingDescriptor.NO_RESCALE;
}
@@ -337,7 +340,8 @@ class TaskStateAssignment {
assignment -> mappingRetriever.apply(assignment, true),
subtaskGateOrPartitionMappings,
subtaskMappingCalculator,
- rescaledChannelsMappings);
+ rescaledChannelsMappings,
+ isInput);
if (Arrays.stream(gateOrPartitionDescriptors)
.allMatch(InflightDataGateOrPartitionRescalingDescriptor::isIdentity)) {
@@ -356,10 +360,14 @@ class TaskStateAssignment {
Function<TaskStateAssignment, SubtasksRescaleMapping> mappingCalculator,
Map<Integer, SubtasksRescaleMapping> subtaskGateOrPartitionMappings,
Function<Integer, SubtasksRescaleMapping> subtaskMappingCalculator,
- SubtasksRescaleMapping[] rescaledChannelsMappings) {
+ SubtasksRescaleMapping[] rescaledChannelsMappings,
+ boolean isInput) {
return IntStream.range(0, rescaledChannelsMappings.length)
.mapToObj(
partition -> {
+ if (!hasInFlightData(isInput, partition)) {
+ return InflightDataGateOrPartitionRescalingDescriptor.NO_STATE;
+ }
TaskStateAssignment connectedAssignment =
connectedAssignments[partition];
SubtasksRescaleMapping rescaleMapping =
@@ -381,6 +389,14 @@ class TaskStateAssignment {
.toArray(InflightDataGateOrPartitionRescalingDescriptor[]::new);
}
+ private boolean hasInFlightData(boolean isInput, int gateOrPartitionIndex) {
+ if (isInput) {
+ return hasInFlightDataForInputGate(gateOrPartitionIndex);
+ } else {
+ return hasInFlightDataForResultPartition(gateOrPartitionIndex);
+ }
+ }
+
private InflightDataGateOrPartitionRescalingDescriptor
getInflightDataGateOrPartitionRescalingDescriptor(
OperatorInstanceID instanceID,
@@ -479,6 +495,51 @@ class TaskStateAssignment {
checkSubtaskMapping(oldMapping, mapping, mapper.isAmbiguous()));
}
+ public boolean hasInFlightDataForInputGate(int gateIndex) {
+ // Check own input state for this gate
+ if (inputStateGates.contains(gateIndex)) {
+ return true;
+ }
+
+ // Check upstream output state for this gate
+ TaskStateAssignment upstreamAssignment = getUpstreamAssignments()[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 false;
+ }
+
+ public boolean hasInFlightDataForResultPartition(int partitionIndex) {
+ // Check own output state for this partition
+ if (outputStatePartitions.contains(partitionIndex)) {
+ return true;
+ }
+
+ // Check downstream input state for this partition
+ TaskStateAssignment downstreamAssignment = getDownstreamAssignments()[partitionIndex];
+
+ if (downstreamAssignment != null && downstreamAssignment.hasInputState()) {
+ IntermediateResult producedResult =
+ executionJobVertex.getProducedDataSets()[partitionIndex];
+ IntermediateDataSetID resultId = producedResult.getId();
+ List<IntermediateResult> 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 false;
+ }
+
@Override
public String toString() {
return "TaskStateAssignment for " + executionJobVertex.getName();
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/InflightDataRescalingDescriptorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/InflightDataRescalingDescriptorTest.java
new file mode 100644
index 00000000000..6252d5b2106
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/InflightDataRescalingDescriptorTest.java
@@ -0,0 +1,121 @@
+/*
+ * 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.runtime.checkpoint;
+
+import org.apache.flink.runtime.checkpoint.InflightDataRescalingDescriptor.InflightDataGateOrPartitionRescalingDescriptor;
+import org.apache.flink.runtime.checkpoint.InflightDataRescalingDescriptor.InflightDataGateOrPartitionRescalingDescriptor.MappingType;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Arrays;
+import java.util.Collections;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link InflightDataRescalingDescriptor}. */
+class InflightDataRescalingDescriptorTest {
+
+ @Test
+ void testNoStateDescriptorThrowsOnGetOldSubtaskInstances() {
+ InflightDataGateOrPartitionRescalingDescriptor noStateDescriptor =
+ InflightDataGateOrPartitionRescalingDescriptor.NO_STATE;
+
+ assertThatThrownBy(noStateDescriptor::getOldSubtaskInstances)
+ .isInstanceOf(UnsupportedOperationException.class)
+ .hasMessageContaining(
+ "Cannot get old subtasks from a descriptor that represents no state");
+ }
+
+ @Test
+ void testNoStateDescriptorThrowsOnGetRescaleMappings() {
+ InflightDataGateOrPartitionRescalingDescriptor noStateDescriptor =
+ InflightDataGateOrPartitionRescalingDescriptor.NO_STATE;
+
+ assertThatThrownBy(noStateDescriptor::getRescaleMappings)
+ .isInstanceOf(UnsupportedOperationException.class)
+ .hasMessageContaining(
+ "Cannot get rescale mappings from a descriptor that represents no state");
+ }
+
+ @Test
+ void testNoStateDescriptorIsIdentity() {
+ InflightDataGateOrPartitionRescalingDescriptor noStateDescriptor =
+ InflightDataGateOrPartitionRescalingDescriptor.NO_STATE;
+
+ assertThat(noStateDescriptor.isIdentity()).isTrue();
+ }
+
+ @Test
+ void testRegularDescriptorDoesNotThrow() {
+ int[] oldSubtasks = new int[] {0, 1, 2};
+ RescaleMappings mappings =
+ RescaleMappings.of(Arrays.stream(new int[][] {{0}, {1}, {2}}), 3);
+
+ InflightDataGateOrPartitionRescalingDescriptor descriptor =
+ new InflightDataGateOrPartitionRescalingDescriptor(
+ oldSubtasks, mappings, Collections.emptySet(), MappingType.RESCALING);
+
+ // Should not throw
+ assertThat(descriptor.getOldSubtaskInstances()).isEqualTo(oldSubtasks);
+ assertThat(descriptor.getRescaleMappings()).isEqualTo(mappings);
+ assertThat(descriptor.isIdentity()).isFalse();
+ }
+
+ @Test
+ void testIdentityDescriptor() {
+ int[] oldSubtasks = new int[] {0};
+ RescaleMappings mappings = RescaleMappings.identity(1, 1);
+
+ InflightDataGateOrPartitionRescalingDescriptor descriptor =
+ new InflightDataGateOrPartitionRescalingDescriptor(
+ oldSubtasks, mappings, Collections.emptySet(), MappingType.IDENTITY);
+
+ assertThat(descriptor.isIdentity()).isTrue();
+ assertThat(descriptor.getOldSubtaskInstances()).isEqualTo(oldSubtasks);
+ assertThat(descriptor.getRescaleMappings()).isEqualTo(mappings);
+ }
+
+ @Test
+ void testInflightDataRescalingDescriptorWithNoStateDescriptor() {
+ // Create a descriptor array with NO_STATE descriptor
+ InflightDataGateOrPartitionRescalingDescriptor[] descriptors =
+ new InflightDataGateOrPartitionRescalingDescriptor[] {
+ InflightDataGateOrPartitionRescalingDescriptor.NO_STATE,
+ new InflightDataGateOrPartitionRescalingDescriptor(
+ new int[] {0, 1},
+ RescaleMappings.of(Arrays.stream(new int[][] {{0}, {1}}), 2),
+ Collections.emptySet(),
+ MappingType.RESCALING)
+ };
+
+ InflightDataRescalingDescriptor rescalingDescriptor =
+ new InflightDataRescalingDescriptor(descriptors);
+
+ // First gate/partition has NO_STATE
+ assertThatThrownBy(() -> rescalingDescriptor.getOldSubtaskIndexes(0))
+ .isInstanceOf(UnsupportedOperationException.class);
+ assertThatThrownBy(() -> rescalingDescriptor.getChannelMapping(0))
+ .isInstanceOf(UnsupportedOperationException.class);
+
+ // Second gate/partition has normal state
+ assertThat(rescalingDescriptor.getOldSubtaskIndexes(1)).isEqualTo(new int[] {0, 1});
+ assertThat(rescalingDescriptor.getChannelMapping(1)).isNotNull();
+ }
+}
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 97b8afbdd9f..fc393090d5b 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
@@ -509,7 +509,7 @@ class StateAssignmentOperationTest {
Stream.of(upstream1, upstream2, downstream)
.map(v -> v.getOperatorIDs().get(0).getGeneratedOperatorID())
.collect(Collectors.toList());
- Map<OperatorID, OperatorState> states = buildOperatorStates(operatorIds, 3);
+ Map<OperatorID, OperatorState> states = buildOperatorStatesForTwoGates(operatorIds, 3);
connectVertices(upstream1, downstream, ARBITRARY, RANGE);
connectVertices(upstream2, downstream, ROUND_ROBIN, ROUND_ROBIN);
@@ -1135,6 +1135,89 @@ class StateAssignmentOperationTest {
}));
}
+ private Map<OperatorID, OperatorState> buildOperatorStatesForTwoGates(
+ List<OperatorID> operatorIDs, int numSubTasks) {
+ Random random = new Random();
+ // operatorIDs should be [upstream1, upstream2, downstream]
+ // downstream has 2 input gates (from upstream1 and upstream2)
+ return operatorIDs.stream()
+ .collect(
+ Collectors.toMap(
+ Function.identity(),
+ operatorID -> {
+ OperatorState state =
+ new OperatorState(
+ "", "", operatorID, numSubTasks, MAX_P);
+ for (int i = 0; i < numSubTasks; i++) {
+ OperatorSubtaskState.Builder builder =
+ OperatorSubtaskState.builder()
+ .setManagedOperatorState(
+ new StateObjectCollection<>(
+ asList(
+ createNewOperatorStateHandle(
+ 10, random),
+ createNewOperatorStateHandle(
+ 10,
+ random))))
+ .setRawOperatorState(
+ new StateObjectCollection<>(
+ asList(
+ createNewOperatorStateHandle(
+ 5, random),
+ createNewOperatorStateHandle(
+ 5,
+ random))))
+ .setManagedKeyedState(
+ StateObjectCollection.singleton(
+ createNewKeyedStateHandle(
+ KeyGroupRange.of(
+ i, i))))
+ .setRawKeyedState(
+ StateObjectCollection.singleton(
+ createNewKeyedStateHandle(
+ KeyGroupRange.of(
+ i, i))));
+
+ // Handle input channel state
+ if (operatorID == operatorIDs.get(2)) {
+ // This is the downstream operator with 2 input gates
+ builder.setInputChannelState(
+ new StateObjectCollection<>(
+ asList(
+ createNewInputChannelStateHandle(
+ 10, 0,
+ random), // gate 0
+ createNewInputChannelStateHandle(
+ 10, 1, random) // gate 1
+ )));
+ } else {
+ // Upstream operators don't have input state
+ builder.setInputChannelState(
+ StateObjectCollection.empty());
+ }
+
+ // Handle result subpartition state
+ if (operatorID != operatorIDs.get(2)) {
+ // Upstream operators have output state
+ builder.setResultSubpartitionState(
+ new StateObjectCollection<>(
+ asList(
+ createNewResultSubpartitionStateHandle(
+ 10, 0, random),
+ createNewResultSubpartitionStateHandle(
+ 10, 0, random))));
+ } else {
+ // Downstream operator doesn't have output state
+ builder.setResultSubpartitionState(
+ StateObjectCollection.empty());
+ }
+
+ state.putState(i, builder.build());
+ }
+ return state;
+ }));
+ }
+
private static class OperatorIdWithParallelism {
private final OperatorID operatorID;
private final int parallelism;
@@ -1276,4 +1359,217 @@ class StateAssignmentOperationTest {
.getTaskStateSnapshot()
.getSubtaskStateByOperatorID(operatorId);
}
+
+ @Test
+ void testMixedExchangesForwardAndHashNoStateOnForward()
+ throws JobException, JobExecutionException {
+ // Create topology: source -> (forward to map1, hash to map2)
+ JobVertex source = createJobVertex(new OperatorID(), 2);
+ JobVertex map1 = createJobVertex(new OperatorID(), 2);
+ JobVertex map2 = createJobVertex(new OperatorID(), 3);
+
+ List<OperatorID> operatorIds =
+ Stream.of(source, map1, map2)
+ .map(v -> v.getOperatorIDs().get(0).getGeneratedOperatorID())
+ .collect(Collectors.toList());
+
+ // Create state with output state only for hash exchange (to map2)
+ Map<OperatorID, OperatorState> states = new HashMap<>();
+ Random random = new Random();
+
+ // Source has output state only for partition 1 (hash exchange)
+ OperatorState sourceState = new OperatorState("", "", operatorIds.get(0), 2, MAX_P);
+ for (int i = 0; i < 2; i++) {
+ sourceState.putState(
+ i,
+ OperatorSubtaskState.builder()
+ .setResultSubpartitionState(
+ new StateObjectCollection<>(
+ Arrays.asList(
+ // No state for partition 0 (forward)
+ createNewResultSubpartitionStateHandle(
+ 10, 1, random) // partition 1 (hash)
+ )))
+ .build());
+ }
+ states.put(operatorIds.get(0), sourceState);
+
+ // Map1 (forward) has no input state
+ OperatorState map1State = new OperatorState("", "", operatorIds.get(1), 2, MAX_P);
+ for (int i = 0; i < 2; i++) {
+ map1State.putState(i, OperatorSubtaskState.builder().build());
+ }
+ states.put(operatorIds.get(1), map1State);
+
+ // Map2 (hash) has input state
+ OperatorState map2State = new OperatorState("", "", operatorIds.get(2), 2, MAX_P);
+ for (int i = 0; i < 2; i++) {
+ map2State.putState(
+ i,
+ OperatorSubtaskState.builder()
+ .setInputChannelState(
+ new StateObjectCollection<>(
+ Arrays.asList(
+ createNewInputChannelStateHandle(
+ 10, 0, random))))
+ .build());
+ }
+ states.put(operatorIds.get(2), map2State);
+
+ // Connect vertices
+ connectVertices(source, map1, RANGE, RANGE); // Forward-like connection
+ connectVertices(source, map2, ARBITRARY, RANGE); // Hash connection
+
+ Map<OperatorID, ExecutionJobVertex> vertices = toExecutionVertices(source, map1, map2);
+
+ // This should not throw UnsupportedOperationException
+ new StateAssignmentOperation(0, new HashSet<>(vertices.values()), states, false)
+ .assignStates();
+
+ // Verify state assignment succeeded
+ assertThat(getAssignedState(vertices.get(operatorIds.get(2)), operatorIds.get(2), 0))
+ .isNotNull();
+ }
+
+ @Test
+ void testMixedExchangesMultipleGatesWithPartialState()
+ throws JobException, JobExecutionException {
+ // Create topology with 3 upstreams connecting to 1 downstream
+ JobVertex upstream1 = createJobVertex(new OperatorID(), 2);
+ JobVertex upstream2 = createJobVertex(new OperatorID(), 2);
+ JobVertex upstream3 = createJobVertex(new OperatorID(), 2);
+ JobVertex downstream = createJobVertex(new OperatorID(), 2);
+
+ List<OperatorID> operatorIds =
+ Stream.of(upstream1, upstream2, upstream3, downstream)
+ .map(v -> v.getOperatorIDs().get(0).getGeneratedOperatorID())
+ .collect(Collectors.toList());
+
+ // Build state where only upstream2 has output state
+ Map<OperatorID, OperatorState> states = new HashMap<>();
+ Random random = new Random();
+
+ // Upstream1 - no output state
+ OperatorState upstream1State = new OperatorState("", "", operatorIds.get(0), 3, MAX_P);
+ for (int i = 0; i < 3; i++) {
+ upstream1State.putState(i, OperatorSubtaskState.builder().build());
+ }
+ states.put(operatorIds.get(0), upstream1State);
+
+ // Upstream2 - has output state
+ OperatorState upstream2State = new OperatorState("", "", operatorIds.get(1), 3, MAX_P);
+ for (int i = 0; i < 3; i++) {
+ upstream2State.putState(
+ i,
+ OperatorSubtaskState.builder()
+ .setResultSubpartitionState(
+ new StateObjectCollection<>(
+ Arrays.asList(
+ createNewResultSubpartitionStateHandle(
+ 10, 0, random))))
+ .build());
+ }
+ states.put(operatorIds.get(1), upstream2State);
+
+ // Upstream3 - no output state
+ OperatorState upstream3State = new OperatorState("", "", operatorIds.get(2), 3, MAX_P);
+ for (int i = 0; i < 3; i++) {
+ upstream3State.putState(i, OperatorSubtaskState.builder().build());
+ }
+ states.put(operatorIds.get(2), upstream3State);
+
+ // Downstream - has input state only for gate 1 (from upstream2)
+ OperatorState downstreamState = new OperatorState("", "", operatorIds.get(3), 3, MAX_P);
+ for (int i = 0; i < 3; i++) {
+ downstreamState.putState(
+ i,
+ OperatorSubtaskState.builder()
+ .setInputChannelState(
+ new StateObjectCollection<>(
+ Arrays.asList(
+ createNewInputChannelStateHandle(
+ 10, 1, random) // gate 1 only
+ )))
+ .build());
+ }
+ states.put(operatorIds.get(3), downstreamState);
+
+ // Connect all upstreams to downstream
+ connectVertices(upstream1, downstream, RANGE, RANGE); // gate 0
+ connectVertices(upstream2, downstream, ARBITRARY, RANGE); // gate 1
+ connectVertices(upstream3, downstream, ROUND_ROBIN, ROUND_ROBIN); // gate 2
+
+ Map<OperatorID, ExecutionJobVertex> vertices =
+ toExecutionVertices(upstream1, upstream2, upstream3, downstream);
+
+ // This should not throw UnsupportedOperationException
+ new StateAssignmentOperation(0, new HashSet<>(vertices.values()), states, false)
+ .assignStates();
+
+ // Verify downstream received state
+ OperatorSubtaskState downstreamAssignedState =
+ getAssignedState(vertices.get(operatorIds.get(3)), operatorIds.get(3), 0);
+ assertThat(downstreamAssignedState).isNotNull();
+ assertThat(downstreamAssignedState.getInputChannelState()).isNotEmpty();
+ }
+
+ @Test
+ void testMixedExchangesRescaleAndRebalanceNoStateOnRescale()
+ throws JobException, JobExecutionException {
+ // Create topology with mixed partitioner types
+ JobVertex source = createJobVertex(new OperatorID(), 4);
+ JobVertex sink = createJobVertex(new OperatorID(), 2);
+
+ List<OperatorID> operatorIds =
+ Stream.of(source, sink)
+ .map(v -> v.getOperatorIDs().get(0).getGeneratedOperatorID())
+ .collect(Collectors.toList());
+
+ // Create state - source has output state
+ Map<OperatorID, OperatorState> states = new HashMap<>();
+ Random random = new Random();
+
+ OperatorState sourceState = new OperatorState("", "", operatorIds.get(0), 4, MAX_P);
+ for (int i = 0; i < 4; i++) {
+ sourceState.putState(
+ i,
+ OperatorSubtaskState.builder()
+ .setResultSubpartitionState(
+ new StateObjectCollection<>(
+ Arrays.asList(
+ createNewResultSubpartitionStateHandle(
+ 10, 0, random))))
+ .build());
+ }
+ states.put(operatorIds.get(0), sourceState);
+
+ // Sink has input state
+ OperatorState sinkState = new OperatorState("", "", operatorIds.get(1), 4, MAX_P);
+ for (int i = 0; i < 4; i++) {
+ sinkState.putState(
+ i,
+ OperatorSubtaskState.builder()
+ .setInputChannelState(
+ new StateObjectCollection<>(
+ Arrays.asList(
+ createNewInputChannelStateHandle(
+ 10, 0, random))))
+ .build());
+ }
+ states.put(operatorIds.get(1), sinkState);
+
+ // Connect with RESCALE partitioner
+ connectVertices(source, sink, ROUND_ROBIN, ROUND_ROBIN);
+
+ Map<OperatorID, ExecutionJobVertex> vertices = toExecutionVertices(source, sink);
+
+ // This should succeed even with RESCALE partitioner when parallelism changes
+ new StateAssignmentOperation(0, new HashSet<>(vertices.values()), states, false)
+ .assignStates();
+
+ // Verify state was assigned
+ OperatorSubtaskState sinkAssignedState =
+ getAssignedState(vertices.get(operatorIds.get(1)), operatorIds.get(1), 0);
+ assertThat(sinkAssignedState).isNotNull();
+ }
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/StateHandleDummyUtil.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/StateHandleDummyUtil.java
index 60ec98e7e0f..7007821a6fe 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/StateHandleDummyUtil.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/StateHandleDummyUtil.java
@@ -144,16 +144,26 @@ public class StateHandleDummyUtil {
public static InputChannelStateHandle createNewInputChannelStateHandle(
int numNamedStates, Random random) {
+ return createNewInputChannelStateHandle(numNamedStates, 0, random);
+ }
+
+ public static InputChannelStateHandle createNewInputChannelStateHandle(
+ int numNamedStates, int gateIndex, Random random) {
return new InputChannelStateHandle(
- new InputChannelInfo(0, random.nextInt()),
+ new InputChannelInfo(gateIndex, random.nextInt()),
createStreamStateHandle(numNamedStates, random),
genOffsets(numNamedStates, random));
}
public static ResultSubpartitionStateHandle createNewResultSubpartitionStateHandle(
int numNamedStates, Random random) {
+ return createNewResultSubpartitionStateHandle(numNamedStates, 0, random);
+ }
+
+ public static ResultSubpartitionStateHandle createNewResultSubpartitionStateHandle(
+ int numNamedStates, int partitionIndex, Random random) {
return new ResultSubpartitionStateHandle(
- new ResultSubpartitionInfo(random.nextInt(), random.nextInt()),
+ new ResultSubpartitionInfo(partitionIndex, random.nextInt()),
createStreamStateHandle(numNamedStates, random),
genOffsets(numNamedStates, random));
}
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 0720197e911..c87cfde3c60 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
@@ -52,6 +52,7 @@ import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
@@ -341,6 +342,30 @@ public class CommonTestUtils {
});
}
+ /** Wait for one checkpoint with in-flight buffers. */
+ public static String waitForCheckpointWithInflightBuffers(JobID jobID, MiniCluster miniCluster)
+ throws Exception {
+ CompletableFuture<String> checkpointPath = new CompletableFuture<>();
+ waitForCheckpoints(
+ jobID,
+ miniCluster,
+ checkpointStatsSnapshot -> {
+ if (checkpointStatsSnapshot == null) {
+ return false;
+ }
+ CompletedCheckpointStats latestCompletedCheckpoint =
+ checkpointStatsSnapshot.getHistory().getLatestCompletedCheckpoint();
+
+ if (latestCompletedCheckpoint == null
+ || latestCompletedCheckpoint.getPersistedData() == 0) {
+ return false;
+ }
+ checkpointPath.complete(latestCompletedCheckpoint.getExternalPath());
+ return true;
+ });
+ return checkpointPath.get();
+ }
+
/** Wait for (at least) the given number of successful checkpoints. */
public static void waitForCheckpoint(JobID jobID, MiniCluster miniCluster, int numCheckpoints)
throws Exception {
diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleWithMixedExchangesITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleWithMixedExchangesITCase.java
new file mode 100644
index 00000000000..c24f9b5b655
--- /dev/null
+++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleWithMixedExchangesITCase.java
@@ -0,0 +1,361 @@
+/*
+ * 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.eventtime.WatermarkStrategy;
+import org.apache.flink.api.common.typeinfo.Types;
+import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy;
+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.StateRecoveryOptions;
+import org.apache.flink.configuration.TaskManagerOptions;
+import org.apache.flink.connector.datagen.source.DataGeneratorSource;
+import org.apache.flink.core.execution.JobClient;
+import org.apache.flink.runtime.minicluster.MiniCluster;
+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.test.util.MiniClusterWithClientResource;
+import org.apache.flink.util.TestLogger;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.time.Duration;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Random;
+
+/**
+ * Integration test for rescaling jobs with mixed (UC-supported and UC-unsupported) exchanges from
+ * an unaligned checkpoint.
+ */
+@RunWith(Parameterized.class)
+public class UnalignedCheckpointRescaleWithMixedExchangesITCase extends TestLogger {
+
+ private static final int NUM_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 10;
+ private static final int MAX_SLOTS = NUM_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private static final Random RANDOM = new Random();
+
+ private static MiniClusterWithClientResource cluster;
+ @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+ @Parameterized.Parameter public ExecuteJobViaEnv executeJobViaEnv;
+
+ @Parameterized.Parameters(name = "Test case {index}")
+ public static Collection<ExecuteJobViaEnv> parameter() {
+ return List.of(
+ UnalignedCheckpointRescaleWithMixedExchangesITCase::createMultiOutputDAG,
+ UnalignedCheckpointRescaleWithMixedExchangesITCase::createMultiInputDAG,
+ UnalignedCheckpointRescaleWithMixedExchangesITCase::createRescalePartitionerDAG,
+ UnalignedCheckpointRescaleWithMixedExchangesITCase::createMixedComplexityDAG);
+ }
+
+ @Before
+ public void setup() throws Exception {
+ cluster =
+ new MiniClusterWithClientResource(
+ new MiniClusterResourceConfiguration.Builder()
+ .setConfiguration(new Configuration())
+ .setNumberTaskManagers(NUM_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .build());
+ cluster.before();
+ }
+
+ @After
+ public void shutDownExistingCluster() {
+ if (cluster != null) {
+ cluster.after();
+ cluster = null;
+ }
+ }
+
+ /**
+ * Tests rescaling from an unaligned checkpoint with different job structures that have mixed
+ * (UC-supported and UC-unsupported) exchanges.
+ */
+ @Test
+ public void testRescaleFromUnalignedCheckpoint() throws Exception {
+ final MiniCluster miniCluster = cluster.getMiniCluster();
+
+ // Step 1: Run the job with initial parallelism and take a checkpoint
+ JobClient jobClient1 = executeJobViaEnv.executeJob(getUnalignedCheckpointEnv(null));
+
+ CommonTestUtils.waitForJobStatus(jobClient1, Collections.singletonList(JobStatus.RUNNING));
+ CommonTestUtils.waitForAllTaskRunning(miniCluster, jobClient1.getJobID(), false);
+ String checkpointPath =
+ CommonTestUtils.waitForCheckpointWithInflightBuffers(
+ jobClient1.getJobID(), miniCluster);
+ jobClient1.cancel().get();
+
+ // Step 2: Restore the job with a different parallelism
+ JobClient jobClient2 =
+ executeJobViaEnv.executeJob(getUnalignedCheckpointEnv(checkpointPath));
+
+ CommonTestUtils.waitForJobStatus(jobClient2, Collections.singletonList(JobStatus.RUNNING));
+ CommonTestUtils.waitForAllTaskRunning(miniCluster, jobClient2.getJobID(), false);
+ CommonTestUtils.waitForCheckpointWithInflightBuffers(jobClient2.getJobID(), miniCluster);
+ jobClient2.cancel().get();
+ }
+
+ private StreamExecutionEnvironment getUnalignedCheckpointEnv(@Nullable String recoveryPath)
+ throws IOException {
+ Configuration conf = new Configuration();
+ conf.set(CheckpointingOptions.CHECKPOINTING_INTERVAL, Duration.ofSeconds(1));
+ // Disable aligned timeout to ensure it works with unaligned checkpoint directly
+ conf.set(CheckpointingOptions.ALIGNED_CHECKPOINT_TIMEOUT, Duration.ofSeconds(0));
+ conf.set(
+ CheckpointingOptions.EXTERNALIZED_CHECKPOINT_RETENTION,
+ ExternalizedCheckpointRetention.RETAIN_ON_CANCELLATION);
+ conf.set(
+ CheckpointingOptions.CHECKPOINTS_DIRECTORY,
+ temporaryFolder.newFolder().toURI().toString());
+ conf.set(CheckpointingOptions.ENABLE_UNALIGNED, true);
+ // Decrease the memory segment size to avoid the test is so slow for some reasons:
+ // 1. When a flink job recovers from unaligned checkpoint, it has to consume all inflight
+ // buffers during recovery phase. The smaller the buffer size, the fewer records are
+ // snapshotted during the checkpoint, resulting in fewer records are needed to be consumed
+ // during recovery.
+ // 2. Forward or rescale exchange does not support unaligned checkpoint, it means Forward
+ // or rescale exchanges are still using aligned checkpoint even if unaligned checkpoint is
+ // enabled. All buffers(records) before barrier must be consumed for aligned checkpoint.
+ // The smaller the buffer size means the fewer records are needed to be consumed during
+ // aligned checkpoint.
+ conf.set(TaskManagerOptions.MEMORY_SEGMENT_SIZE, MemorySize.parse("1 kb"));
+ // To prevent the picked checkpoint is deleted
+ conf.set(CheckpointingOptions.MAX_RETAINED_CHECKPOINTS, 50);
+ if (recoveryPath != null) {
+ conf.set(StateRecoveryOptions.SAVEPOINT_PATH, recoveryPath);
+ }
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(conf);
+ env.disableOperatorChaining();
+ return env;
+ }
+
+ private static JobClient createMultiOutputDAG(StreamExecutionEnvironment env) throws Exception {
+ DataGeneratorSource<Long> source =
+ new DataGeneratorSource<>(
+ index -> index,
+ Long.MAX_VALUE,
+ RateLimiterStrategy.perSecond(5000),
+ Types.LONG);
+
+ int sourceParallelism = getRandomParallelism();
+ DataStream<Long> sourceStream =
+ env.fromSource(source, WatermarkStrategy.noWatermarks(), "Data Generator")
+ .setParallelism(sourceParallelism);
+
+ sourceStream
+ .keyBy((KeySelector<Long, Long>) value -> value)
+ .map(
+ x -> {
+ Thread.sleep(5);
+ return x;
+ })
+ .name("Map after keyBy")
+ .setParallelism(getRandomParallelism());
+
+ // Keep the same parallelism to ensure the ForwardPartitioner will be used.
+ sourceStream
+ .map(
+ x -> {
+ Thread.sleep(1);
+ return x;
+ })
+ .name("Map after forward")
+ .setParallelism(sourceParallelism);
+
+ return env.executeAsync();
+ }
+
+ private static JobClient createMultiInputDAG(StreamExecutionEnvironment env) throws Exception {
+ int source1Parallelism = getRandomParallelism();
+ DataGeneratorSource<Long> source1 =
+ new DataGeneratorSource<>(
+ index -> index,
+ Long.MAX_VALUE,
+ RateLimiterStrategy.perSecond(5000),
+ Types.LONG);
+ DataStream<Long> sourceStream1 =
+ env.fromSource(source1, WatermarkStrategy.noWatermarks(), "Source 1")
+ .setParallelism(source1Parallelism);
+
+ int source2Parallelism = getRandomParallelism();
+ DataGeneratorSource<Long> source2 =
+ new DataGeneratorSource<>(
+ index -> index,
+ Long.MAX_VALUE,
+ RateLimiterStrategy.perSecond(5000),
+ Types.LONG);
+ DataStream<Long> sourceStream2 =
+ env.fromSource(source2, WatermarkStrategy.noWatermarks(), "Source 2")
+ .setParallelism(source2Parallelism);
+
+ // Keep the same parallelism to ensure the ForwardPartitioner will be used.
+ DataStream<Long> forwardedStream =
+ sourceStream2.map(x -> x).setParallelism(source2Parallelism);
+
+ sourceStream1
+ .rebalance()
+ .connect(forwardedStream)
+ .map(new SleepingCoMap())
+ .name("Co-Map")
+ .setParallelism(getRandomParallelism());
+
+ return env.executeAsync();
+ }
+
+ private static JobClient createRescalePartitionerDAG(StreamExecutionEnvironment env)
+ throws Exception {
+ DataGeneratorSource<Long> source =
+ new DataGeneratorSource<>(
+ index -> index,
+ Long.MAX_VALUE,
+ RateLimiterStrategy.perSecond(5000),
+ Types.LONG);
+
+ int sourceParallelism = getRandomParallelism();
+ DataStream<Long> sourceStream =
+ env.fromSource(source, WatermarkStrategy.noWatermarks(), "Data Generator")
+ .setParallelism(sourceParallelism);
+
+ sourceStream
+ .keyBy((KeySelector<Long, Long>) value -> value)
+ .map(
+ x -> {
+ Thread.sleep(5);
+ return x;
+ })
+ .name("Map after keyBy")
+ .setParallelism(getRandomParallelism());
+
+ sourceStream
+ .rescale()
+ .map(
+ x -> {
+ Thread.sleep(1);
+ return x;
+ })
+ .name("Map after rescale")
+ .setParallelism(getRandomParallelism());
+
+ return env.executeAsync();
+ }
+
+ private static JobClient createMixedComplexityDAG(StreamExecutionEnvironment env)
+ throws Exception {
+ // Multi-input part
+ int source1Parallelism = getRandomParallelism();
+ DataGeneratorSource<Long> source1 =
+ new DataGeneratorSource<>(
+ index -> index,
+ Long.MAX_VALUE,
+ RateLimiterStrategy.perSecond(5000),
+ Types.LONG);
+ DataStream<Long> sourceStream1 =
+ env.fromSource(source1, WatermarkStrategy.noWatermarks(), "Source 1")
+ .setParallelism(source1Parallelism);
+
+ int source2Parallelism = getRandomParallelism();
+ DataGeneratorSource<Long> source2 =
+ new DataGeneratorSource<>(
+ index -> index,
+ Long.MAX_VALUE,
+ RateLimiterStrategy.perSecond(5000),
+ Types.LONG);
+ DataStream<Long> sourceStream2 =
+ env.fromSource(source2, WatermarkStrategy.noWatermarks(), "Source 2")
+ .setParallelism(source2Parallelism);
+
+ // Keep the same parallelism to ensure the ForwardPartitioner will be used.
+ DataStream<Long> forwardedStream =
+ sourceStream2.map(x -> x).setParallelism(source2Parallelism);
+
+ DataStream<Long> multiInputMap =
+ sourceStream1
+ .rebalance()
+ .connect(forwardedStream)
+ .map(new SleepingCoMap())
+ .name("Co-Map")
+ .setParallelism(getRandomParallelism());
+
+ // Multi-output part
+ multiInputMap
+ .keyBy((KeySelector<Long, Long>) value -> value)
+ .map(
+ x -> {
+ Thread.sleep(5);
+ return x;
+ })
+ .name("Map after keyBy")
+ .setParallelism(getRandomParallelism());
+
+ // Keep the same parallelism to ensure the ForwardPartitioner will be used.
+ multiInputMap
+ .map(
+ x -> {
+ Thread.sleep(1);
+ return x;
+ })
+ .name("Map after forward")
+ .setParallelism(multiInputMap.getParallelism());
+
+ return env.executeAsync();
+ }
+
+ private static int getRandomParallelism() {
+ return RANDOM.nextInt(MAX_SLOTS) + 1;
+ }
+
+ /** A simple CoMapFunction that sleeps for 1ms for each element. */
+ private static class SleepingCoMap implements CoMapFunction<Long, Long, Long> {
+ @Override
+ public Long map1(Long value) throws Exception {
+ Thread.sleep(1);
+ return value;
+ }
+
+ @Override
+ public Long map2(Long value) throws Exception {
+ Thread.sleep(1);
+ return value;
+ }
+ }
+
+ public interface ExecuteJobViaEnv {
+ JobClient executeJob(StreamExecutionEnvironment env) throws Exception;
+ }
+}