Address Gabor's comments

Author: 1996fanrui

flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/InflightDataRescalingDescriptor.java

@@ -20,6 +20,7 @@
 import java.io.ObjectStreamException;
 import java.io.Serializable;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.Objects;
 import java.util.Set;
 
@@ -116,7 +117,7 @@ public static class InflightDataGateOrPartitionRescalingDescriptor implements Se
                 new InflightDataGateOrPartitionRescalingDescriptor(
                         new int[0],
                         RescaleMappings.identity(0, 0),
-                        java.util.Collections.emptySet(),
+                        Collections.emptySet(),
                         MappingType.IDENTITY) {
 
                     private static final long serialVersionUID = 1L;

flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskStateAssignment.java

@@ -505,12 +505,12 @@ public boolean hasInFlightDataForInputGate(int gateIndex) {
         TaskStateAssignment upstreamAssignment = getUpstreamAssignments()[gateIndex];
         if (upstreamAssignment != null && upstreamAssignment.hasOutputState()) {
             IntermediateResult inputResult = executionJobVertex.getInputs().get(gateIndex);
-            int partitionIndex =
-                    Arrays.asList(inputResult.getProducer().getProducedDataSets())
-                            .indexOf(inputResult);
-
-            if (partitionIndex != -1) {
-                return upstreamAssignment.outputStatePartitions.contains(partitionIndex);
+            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);
+                }
             }
         }
 
@@ -529,11 +529,12 @@ public boolean hasInFlightDataForResultPartition(int partitionIndex) {
         if (downstreamAssignment != null && downstreamAssignment.hasInputState()) {
             IntermediateResult producedResult =
                     executionJobVertex.getProducedDataSets()[partitionIndex];
-            int gateIndex =
-                    downstreamAssignment.executionJobVertex.getInputs().indexOf(producedResult);
-
-            if (gateIndex != -1) {
-                return downstreamAssignment.inputStateGates.contains(gateIndex);
+            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;

flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleWithMixedExchangesITCase.java

@@ -143,8 +143,15 @@ private StreamExecutionEnvironment getUnalignedCheckpointEnv(@Nullable String re
                 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. Recovery phase needs to consume all inflight buffers
-        // 2. Forward or rescale exchange does not support unaligned checkpoint.
+        // 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);