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:15 UTC
(flink) branch release-2.1 updated (d67c45a750e -> 90762ab27d0)
This is an automated email from the ASF dual-hosted git repository.
fanrui pushed a change to branch release-2.1
in repository https://gitbox.apache.org/repos/asf/flink.git
from d67c45a750e [FLINK-38286][table] Fix: MAP function with duplicate keys produces non-deterministic results
new eab7e917b0c [FLINK-38267][checkpoint] Refactor hasInputState and hasOutputState related logic in TaskStateAssignment
new 4f17abd8715 [FLINK-38267][checkpoint] Only call channel state rescaling logic for exchange with channel state to avoid UnsupportedOperationException
new 90762ab27d0 [FLINK-38267][checkpoint] Fix the test timeout for UnalignedCheckpointRescaleWithMixedExchangesITCase.testRescaleFromUnalignedCheckpoint
The 3 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.../InflightDataRescalingDescriptor.java | 35 +-
.../checkpoint/StateAssignmentOperation.java | 10 +-
.../runtime/checkpoint/TaskStateAssignment.java | 152 ++++++++-
.../InflightDataRescalingDescriptorTest.java | 121 +++++++
.../checkpoint/StateAssignmentOperationTest.java | 298 ++++++++++++++++-
.../runtime/checkpoint/StateHandleDummyUtil.java | 14 +-
.../flink/runtime/testutils/CommonTestUtils.java | 25 ++
...dCheckpointRescaleWithMixedExchangesITCase.java | 361 +++++++++++++++++++++
8 files changed, 994 insertions(+), 22 deletions(-)
create mode 100644 flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/InflightDataRescalingDescriptorTest.java
create mode 100644 flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleWithMixedExchangesITCase.java
(flink) 03/03: [FLINK-38267][checkpoint] Fix the test timeout for UnalignedCheckpointRescaleWithMixedExchangesITCase.testRescaleFromUnalignedCheckpoint
Posted by fa...@apache.org.
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 90762ab27d07c0339f5c476ec6d7331beab81e08
Author: Rui Fan <fa...@apache.org>
AuthorDate: Mon Aug 25 17:53:51 2025 +0200
[FLINK-38267][checkpoint] Fix the test timeout for UnalignedCheckpointRescaleWithMixedExchangesITCase.testRescaleFromUnalignedCheckpoint
When one task has multiple inputs, and the unaligned checkpoint will be disabled for all inputs once one input exchange does not support unaligned checkpoint. It caused no inflight buffers, but UnalignedCheckpointRescaleWithMixedExchangesITCase.testRescaleFromUnalignedCheckpoint always wait for checkpoint with inflight buffers.
Explicitly specifying rebalance can avoid the forward exchange.
---
.../UnalignedCheckpointRescaleWithMixedExchangesITCase.java | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
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
index c24f9b5b655..34c3277c4a1 100644
--- 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
@@ -230,7 +230,7 @@ public class UnalignedCheckpointRescaleWithMixedExchangesITCase extends TestLogg
sourceStream1
.rebalance()
- .connect(forwardedStream)
+ .connect(forwardedStream.rebalance())
.map(new SleepingCoMap())
.name("Co-Map")
.setParallelism(getRandomParallelism());
@@ -307,7 +307,7 @@ public class UnalignedCheckpointRescaleWithMixedExchangesITCase extends TestLogg
DataStream<Long> multiInputMap =
sourceStream1
.rebalance()
- .connect(forwardedStream)
+ .connect(forwardedStream.rebalance())
.map(new SleepingCoMap())
.name("Co-Map")
.setParallelism(getRandomParallelism());
(flink) 01/03: [FLINK-38267][checkpoint] Refactor hasInputState and hasOutputState related logic in TaskStateAssignment
Posted by fa...@apache.org.
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 eab7e917b0cc0a23d50698db194e5de069ab165b
Author: Rui Fan <19...@gmail.com>
AuthorDate: Wed Aug 20 12:35:49 2025 +0200
[FLINK-38267][checkpoint] Refactor hasInputState and hasOutputState related logic in TaskStateAssignment
---
.../checkpoint/StateAssignmentOperation.java | 4 +-
.../runtime/checkpoint/TaskStateAssignment.java | 81 +++++++++++++++++++---
2 files changed, 73 insertions(+), 12 deletions(-)
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 6ce8d6a5932..7b41b4981f0 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
@@ -361,7 +361,7 @@ public class StateAssignmentOperation {
public void reDistributeResultSubpartitionStates(TaskStateAssignment assignment) {
// FLINK-31963: We can skip this phase if there is no output state AND downstream has no
// input states
- if (!assignment.hasOutputState && !assignment.hasDownstreamInputStates()) {
+ if (!assignment.hasOutputState() && !assignment.hasDownstreamInputStates()) {
return;
}
@@ -410,7 +410,7 @@ public class StateAssignmentOperation {
public void reDistributeInputChannelStates(TaskStateAssignment stateAssignment) {
// FLINK-31963: We can skip this phase only if there is no input state AND upstream has no
// output states
- if (!stateAssignment.hasInputState && !stateAssignment.hasUpstreamOutputStates()) {
+ if (!stateAssignment.hasInputState() && !stateAssignment.hasUpstreamOutputStates()) {
return;
}
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 75e11d147a4..54810ab7b55 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
@@ -20,6 +20,8 @@ package org.apache.flink.runtime.checkpoint;
import org.apache.flink.runtime.OperatorIDPair;
import org.apache.flink.runtime.checkpoint.InflightDataRescalingDescriptor.InflightDataGateOrPartitionRescalingDescriptor;
import org.apache.flink.runtime.checkpoint.InflightDataRescalingDescriptor.InflightDataGateOrPartitionRescalingDescriptor.MappingType;
+import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo;
+import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo;
import org.apache.flink.runtime.executiongraph.ExecutionJobVertex;
import org.apache.flink.runtime.executiongraph.IntermediateResult;
import org.apache.flink.runtime.io.network.api.writer.SubtaskStateMapper;
@@ -28,6 +30,8 @@ import org.apache.flink.runtime.jobgraph.OperatorID;
import org.apache.flink.runtime.jobgraph.OperatorInstanceID;
import org.apache.flink.runtime.state.InputChannelStateHandle;
import org.apache.flink.runtime.state.KeyedStateHandle;
+import org.apache.flink.runtime.state.MergedInputChannelStateHandle;
+import org.apache.flink.runtime.state.MergedResultSubpartitionStateHandle;
import org.apache.flink.runtime.state.OperatorStateHandle;
import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
import org.apache.flink.runtime.state.StateObject;
@@ -40,6 +44,7 @@ import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.util.Arrays;
+import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -48,6 +53,7 @@ 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;
import static java.util.Collections.emptySet;
@@ -67,12 +73,16 @@ class TaskStateAssignment {
final Map<OperatorID, OperatorState> oldState;
final boolean hasNonFinishedState;
final boolean isFullyFinished;
- final boolean hasInputState;
- final boolean hasOutputState;
final int newParallelism;
final OperatorID inputOperatorID;
final OperatorID outputOperatorID;
+ /** The InputGate set that containing input buffer state. */
+ private final Set<Integer> inputStateGates;
+
+ /** The ResultPartition set that containing input buffer state. */
+ private final Set<Integer> outputStatePartitions;
+
final Map<OperatorInstanceID, List<OperatorStateHandle>> subManagedOperatorState;
final Map<OperatorInstanceID, List<OperatorStateHandle>> subRawOperatorState;
final Map<OperatorInstanceID, List<KeyedStateHandle>> subManagedKeyedState;
@@ -131,12 +141,63 @@ class TaskStateAssignment {
outputOperatorID = operatorIDs.get(0).getGeneratedOperatorID();
inputOperatorID = operatorIDs.get(operatorIDs.size() - 1).getGeneratedOperatorID();
- hasInputState =
- oldState.get(inputOperatorID).getStates().stream()
- .anyMatch(subState -> !subState.getInputChannelState().isEmpty());
- hasOutputState =
- oldState.get(outputOperatorID).getStates().stream()
- .anyMatch(subState -> !subState.getResultSubpartitionState().isEmpty());
+ inputStateGates = extractInputStateGates(oldState.get(inputOperatorID));
+ outputStatePartitions = extractOutputStatePartitions(oldState.get(outputOperatorID));
+ }
+
+ private static Set<Integer> extractInputStateGates(OperatorState operatorState) {
+ return operatorState.getStates().stream()
+ .map(OperatorSubtaskState::getInputChannelState)
+ .flatMap(Collection::stream)
+ .flatMapToInt(
+ handle -> {
+ if (handle instanceof InputChannelStateHandle) {
+ return IntStream.of(
+ ((InputChannelStateHandle) handle).getInfo().getGateIdx());
+ } else if (handle instanceof MergedInputChannelStateHandle) {
+ return ((MergedInputChannelStateHandle) handle)
+ .getInfos().stream().mapToInt(InputChannelInfo::getGateIdx);
+ } else {
+ throw new IllegalStateException(
+ "Invalid input channel state : " + handle.getClass());
+ }
+ })
+ .distinct()
+ .boxed()
+ .collect(Collectors.toSet());
+ }
+
+ private static Set<Integer> extractOutputStatePartitions(OperatorState operatorState) {
+ return operatorState.getStates().stream()
+ .map(OperatorSubtaskState::getResultSubpartitionState)
+ .flatMap(Collection::stream)
+ .flatMapToInt(
+ handle -> {
+ if (handle instanceof ResultSubpartitionStateHandle) {
+ return IntStream.of(
+ ((ResultSubpartitionStateHandle) handle)
+ .getInfo()
+ .getPartitionIdx());
+ } else if (handle instanceof MergedResultSubpartitionStateHandle) {
+ return ((MergedResultSubpartitionStateHandle) handle)
+ .getInfos().stream()
+ .mapToInt(ResultSubpartitionInfo::getPartitionIdx);
+ } else {
+ throw new IllegalStateException(
+ "Invalid output channel state : " + handle.getClass());
+ }
+ })
+ .distinct()
+ .boxed()
+ .collect(Collectors.toSet());
+ }
+
+ public boolean hasInputState() {
+ return !inputStateGates.isEmpty();
+ }
+
+ public boolean hasOutputState() {
+ return !outputStatePartitions.isEmpty();
}
public TaskStateAssignment[] getDownstreamAssignments() {
@@ -212,7 +273,7 @@ class TaskStateAssignment {
if (hasUpstreamOutputStates == null) {
hasUpstreamOutputStates =
Arrays.stream(getUpstreamAssignments())
- .anyMatch(assignment -> assignment.hasOutputState);
+ .anyMatch(TaskStateAssignment::hasOutputState);
}
return hasUpstreamOutputStates;
}
@@ -221,7 +282,7 @@ class TaskStateAssignment {
if (hasDownstreamInputStates == null) {
hasDownstreamInputStates =
Arrays.stream(getDownstreamAssignments())
- .anyMatch(assignment -> assignment.hasInputState);
+ .anyMatch(TaskStateAssignment::hasInputState);
}
return hasDownstreamInputStates;
}
(flink) 02/03: [FLINK-38267][checkpoint] Only call channel state rescaling logic for exchange with channel state to avoid UnsupportedOperationException
Posted by fa...@apache.org.
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;
+ }
+}