You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by za...@apache.org on 2025/03/04 11:55:22 UTC

(flink) branch master updated (66104a3fbbc -> 0d1b5aa1cdd)

This is an automated email from the ASF dual-hosted git repository.

zakelly pushed a change to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git


    from 66104a3fbbc [hotfix] Bump forst version to 0.1.7
     new be91c9f278f [FLINK-33734][checkpointing] Introduce interface ChannelState,InputStateHandle,OutputStateHandle
     new 32aa7973daa [FLINK-33734][checkpointing] Use interface replace InputChannelStateHandle and ResultSubpartitionStateHandle in OperatorSubtaskState
     new db07fe9f67e [FLINK-33734][checkpointing] Extract interface ChannelStateHandleSerializer
     new 6f6bed67522 [FLINK-33734][checkpointing] Introduce merged channel state handle
     new 88c75d50e8e [hotfix] Make some improper static method become instance method
     new c446f5aae7f [FLINK-33734][checkpointing] Introduce MetadataV6Serializer for merged channel state handle
     new 0d1b5aa1cdd [FLINK-33734][checkpointing] Compact and serialize original channel state handle before send checkpoint ack to JobManager

The 7 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:
 .../flink/runtime/checkpoint/Checkpoints.java      |   4 +-
 .../runtime/checkpoint/OperatorSubtaskState.java   |  37 ++--
 .../PrioritizedOperatorSubtaskState.java           |  99 +++++++--
 .../checkpoint/StateAssignmentOperation.java       |   9 +-
 .../runtime/checkpoint/TaskStateAssignment.java    |  12 +-
 .../channel/SequentialChannelStateReaderImpl.java  |   6 +-
 .../metadata/ChannelStateHandleSerializer.java     | 105 ++-------
 ...er.java => ChannelStateHandleSerializerV1.java} |  69 +++---
 .../metadata/ChannelStateHandleSerializerV2.java   | 169 +++++++++++++++
 .../checkpoint/metadata/MetadataSerializers.java   |   3 +-
 .../metadata/MetadataV2V3SerializerBase.java       |  22 +-
 .../checkpoint/metadata/MetadataV3Serializer.java  |  74 +++----
 .../checkpoint/metadata/MetadataV6Serializer.java  |  30 ++-
 .../state/AbstractMergedChannelStateHandle.java    | 240 +++++++++++++++++++++
 ...ointKeyedStateHandle.java => ChannelState.java} |   4 +-
 .../flink/runtime/state/ChannelStateHelper.java    | 186 ++++++++++++++++
 .../runtime/state/InputChannelStateHandle.java     |   3 +-
 ...KeyedStateHandle.java => InputStateHandle.java} |   4 +-
 .../state/MergedInputChannelStateHandle.java       |  75 +++++++
 .../state/MergedResultSubpartitionStateHandle.java |  76 +++++++
 ...eyedStateHandle.java => OutputStateHandle.java} |   4 +-
 .../state/ResultSubpartitionStateHandle.java       |   2 +-
 .../api/operators/OperatorSnapshotFinalizer.java   |   8 +-
 .../api/operators/OperatorSnapshotFutures.java     |  21 +-
 .../streaming/runtime/tasks/OperatorChain.java     |   6 +-
 .../CheckpointCoordinatorFailureTest.java          |  15 +-
 .../SequentialChannelStateReaderImplTest.java      |   7 +-
 .../ChannelStateHandleSerializerV1Test.java        |  27 ++-
 .../ChannelStateHandleSerializerV2Test.java        |  57 +++++
 .../checkpoint/metadata/ChannelStateTestUtils.java | 175 +++++++++++++++
 .../metadata/MetadataV6SerializerTest.java         | 152 +++++++++++++
 .../runtime/state/ChannelPersistenceITCase.java    |   7 +-
 .../state/MergedChannelStateHandleTest.java        |  49 +++++
 .../operators/OperatorSnapshotFinalizerTest.java   |   6 +-
 .../api/operators/OperatorSnapshotFuturesTest.java |  17 +-
 .../operators/StreamOperatorStateHandlerTest.java  |   8 +-
 .../flink/streaming/util/OperatorSnapshotUtil.java |  52 ++---
 .../runtime/tasks/LocalStateForwardingTest.java    |   9 +-
 38 files changed, 1514 insertions(+), 335 deletions(-)
 copy flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/{ChannelStateHandleSerializer.java => ChannelStateHandleSerializerV1.java} (62%)
 create mode 100644 flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV2.java
 copy flink-core/src/main/java/org/apache/flink/api/common/functions/util/ListCollector.java => flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV6Serializer.java (59%)
 create mode 100644 flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractMergedChannelStateHandle.java
 copy flink-runtime/src/main/java/org/apache/flink/runtime/state/{SavepointKeyedStateHandle.java => ChannelState.java} (84%)
 create mode 100644 flink-runtime/src/main/java/org/apache/flink/runtime/state/ChannelStateHelper.java
 copy flink-runtime/src/main/java/org/apache/flink/runtime/state/{SavepointKeyedStateHandle.java => InputStateHandle.java} (84%)
 create mode 100644 flink-runtime/src/main/java/org/apache/flink/runtime/state/MergedInputChannelStateHandle.java
 create mode 100644 flink-runtime/src/main/java/org/apache/flink/runtime/state/MergedResultSubpartitionStateHandle.java
 copy flink-runtime/src/main/java/org/apache/flink/runtime/state/{SavepointKeyedStateHandle.java => OutputStateHandle.java} (84%)
 copy flink-connectors/flink-hadoop-compatibility/src/test/java/org/apache/flink/hadoopcompatibility/HadoopUtilsTest.java => flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV1Test.java (53%)
 create mode 100644 flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV2Test.java
 create mode 100644 flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateTestUtils.java
 create mode 100644 flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV6SerializerTest.java
 create mode 100644 flink-runtime/src/test/java/org/apache/flink/runtime/state/MergedChannelStateHandleTest.java


(flink) 06/07: [FLINK-33734][checkpointing] Introduce MetadataV6Serializer for merged channel state handle

Posted by za...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

zakelly pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit c446f5aae7f3495eff93233316f78bdb18f0e800
Author: wangfeifan <zo...@163.com>
AuthorDate: Thu Feb 1 18:05:44 2024 +0800

    [FLINK-33734][checkpointing] Introduce MetadataV6Serializer for merged channel state handle
---
 .../checkpoint/metadata/MetadataSerializers.java   |   3 +-
 .../checkpoint/metadata/MetadataV6Serializer.java  |  43 ++++++
 .../metadata/MetadataV6SerializerTest.java         | 152 +++++++++++++++++++++
 3 files changed, 197 insertions(+), 1 deletion(-)

diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataSerializers.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataSerializers.java
index 8646c48e1cf..aa4032683ee 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataSerializers.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataSerializers.java
@@ -30,7 +30,7 @@ import java.util.Map;
 public class MetadataSerializers {
 
     private static final Map<Integer, MetadataSerializer> SERIALIZERS =
-            CollectionUtil.newHashMapWithExpectedSize(5);
+            CollectionUtil.newHashMapWithExpectedSize(6);
 
     static {
         registerSerializer(MetadataV1Serializer.INSTANCE);
@@ -38,6 +38,7 @@ public class MetadataSerializers {
         registerSerializer(MetadataV3Serializer.INSTANCE);
         registerSerializer(MetadataV4Serializer.INSTANCE);
         registerSerializer(MetadataV5Serializer.INSTANCE);
+        registerSerializer(MetadataV6Serializer.INSTANCE);
     }
 
     private static void registerSerializer(MetadataSerializer serializer) {
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV6Serializer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV6Serializer.java
new file mode 100644
index 00000000000..eaa9463cae6
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV6Serializer.java
@@ -0,0 +1,43 @@
+/*
+ * 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.metadata;
+
+import org.apache.flink.annotation.Internal;
+
+/** V6 serializer that use ChannelStateHandleSerializerV2. */
+@Internal
+public class MetadataV6Serializer extends MetadataV5Serializer {
+
+    public static final MetadataSerializer INSTANCE = new MetadataV6Serializer();
+
+    public static final int VERSION = 6;
+
+    private final ChannelStateHandleSerializer channelStateHandleSerializer =
+            new ChannelStateHandleSerializerV2();
+
+    @Override
+    public int getVersion() {
+        return VERSION;
+    }
+
+    @Override
+    protected ChannelStateHandleSerializer getChannelStateHandleSerializer() {
+        return channelStateHandleSerializer;
+    }
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV6SerializerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV6SerializerTest.java
new file mode 100644
index 00000000000..605a47c2794
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV6SerializerTest.java
@@ -0,0 +1,152 @@
+package org.apache.flink.runtime.checkpoint.metadata;
+
+/*
+ * 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.
+ */
+
+import org.apache.flink.core.fs.FileSystem;
+import org.apache.flink.runtime.checkpoint.OperatorState;
+import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
+import org.apache.flink.runtime.checkpoint.StateObjectCollection;
+import org.apache.flink.runtime.state.InputStateHandle;
+import org.apache.flink.runtime.state.OutputStateHandle;
+import org.apache.flink.runtime.state.filesystem.AbstractFsCheckpointStorageAccess;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Random;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
+
+import static java.util.Collections.emptyList;
+import static org.apache.flink.runtime.checkpoint.metadata.ChannelStateTestUtils.randomMergedInputChannelStateHandle;
+import static org.apache.flink.runtime.checkpoint.metadata.ChannelStateTestUtils.randomMergedResultSubpartitionStateHandle;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+/** {@link MetadataV6Serializer} test. */
+class MetadataV6SerializerTest {
+
+    private static final MetadataSerializer INSTANCE = MetadataV6Serializer.INSTANCE;
+
+    private static final Random RND = new Random();
+
+    private String basePath;
+
+    private List<InputStateHandle> inputHandles;
+    private List<OutputStateHandle> outputHandles;
+
+    private CheckpointMetadata metadata;
+
+    @BeforeEach
+    public void beforeEach(@TempDir Path tempDir) throws IOException {
+        basePath = tempDir.toUri().toString();
+
+        final org.apache.flink.core.fs.Path metaPath =
+                new org.apache.flink.core.fs.Path(
+                        basePath, AbstractFsCheckpointStorageAccess.METADATA_FILE_NAME);
+        // this is in the temp folder, so it will get automatically deleted
+        FileSystem.getLocalFileSystem().create(metaPath, FileSystem.WriteMode.OVERWRITE).close();
+    }
+
+    @Test
+    void testSerializeUnmergedChannelStateHandle() throws IOException {
+        testSerializeChannelStateHandle(
+                () ->
+                        ChannelStateTestUtils.randomInputChannelStateHandlesFromSameSubtask()
+                                .stream()
+                                .map(e -> (InputStateHandle) e)
+                                .collect(Collectors.toList()),
+                () ->
+                        ChannelStateTestUtils.randomResultSubpartitionStateHandlesFromSameSubtask()
+                                .stream()
+                                .map(e -> (OutputStateHandle) e)
+                                .collect(Collectors.toList()));
+    }
+
+    @Test
+    void testSerializeMergedChannelStateHandle() throws IOException {
+        testSerializeChannelStateHandle(
+                () -> Collections.singletonList(randomMergedInputChannelStateHandle()),
+                () -> Collections.singletonList(randomMergedResultSubpartitionStateHandle()));
+    }
+
+    private void testSerializeChannelStateHandle(
+            Supplier<List<InputStateHandle>> getter1, Supplier<List<OutputStateHandle>> getter2)
+            throws IOException {
+
+        prepareAndSerializeMetadata(getter1, getter2);
+
+        try (ByteArrayOutputStream out = new ByteArrayOutputStream();
+                DataOutputStream dos = new DataOutputStream(out)) {
+            INSTANCE.serialize(metadata, dos);
+
+            try (DataInputStream dis =
+                    new DataInputStream(new ByteArrayInputStream(out.toByteArray()))) {
+
+                CheckpointMetadata deserializedMetadata =
+                        INSTANCE.deserialize(dis, metadata.getClass().getClassLoader(), basePath);
+
+                Collection<OperatorState> operatorStates = deserializedMetadata.getOperatorStates();
+                assertThat(operatorStates).hasSize(1);
+
+                OperatorState operatorState = operatorStates.iterator().next();
+                assertEquals(1, operatorState.getNumberCollectedStates());
+
+                OperatorSubtaskState subtaskState = operatorState.getState(0);
+
+                assertEquals(inputHandles, subtaskState.getInputChannelState().asList());
+                assertEquals(outputHandles, subtaskState.getResultSubpartitionState().asList());
+            }
+        }
+    }
+
+    private void prepareAndSerializeMetadata(
+            Supplier<List<InputStateHandle>> getter1, Supplier<List<OutputStateHandle>> getter2) {
+        Collection<OperatorState> operatorStates =
+                CheckpointTestUtils.createOperatorStates(RND, basePath, 1, 0, 0, 1);
+
+        inputHandles = getter1.get();
+        outputHandles = getter2.get();
+
+        // Set merged channel state handle to each subtask state
+        for (OperatorState operatorState : operatorStates) {
+            int subtaskStateCount = operatorState.getNumberCollectedStates();
+            for (int i = 0; i < subtaskStateCount; i++) {
+                OperatorSubtaskState originSubtaskState = operatorState.getState(i);
+
+                OperatorSubtaskState.Builder builder = originSubtaskState.toBuilder();
+                builder.setInputChannelState(new StateObjectCollection<>(inputHandles));
+                builder.setResultSubpartitionState(new StateObjectCollection<>(outputHandles));
+
+                operatorState.putState(i, builder.build());
+            }
+        }
+
+        metadata = new CheckpointMetadata(1L, operatorStates, emptyList(), null);
+    }
+}


(flink) 05/07: [hotfix] Make some improper static method become instance method

Posted by za...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

zakelly pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit 88c75d50e8efc0727bbc66dfa2094d7a75d7eb9b
Author: wangfeifan <zo...@163.com>
AuthorDate: Thu Feb 1 17:53:35 2024 +0800

    [hotfix] Make some improper static method become instance method
---
 .../checkpoint/metadata/MetadataV3Serializer.java  | 26 +++++++++++-----------
 .../flink/streaming/util/OperatorSnapshotUtil.java | 22 ++++++++++--------
 2 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV3Serializer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV3Serializer.java
index ed257b9d64d..012914ab14d 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV3Serializer.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV3Serializer.java
@@ -268,39 +268,39 @@ public class MetadataV3Serializer extends MetadataV2V3SerializerBase implements
     }
 
     @VisibleForTesting
-    public static void serializeOperatorStateHandleUtil(
+    public void serializeOperatorStateHandleUtil(
             OperatorStateHandle stateHandle, DataOutputStream dos) throws IOException {
-        INSTANCE.serializeOperatorStateHandle(stateHandle, dos);
+        serializeOperatorStateHandle(stateHandle, dos);
     }
 
     @VisibleForTesting
-    public static OperatorStateHandle deserializeOperatorStateHandleUtil(DataInputStream dis)
+    public OperatorStateHandle deserializeOperatorStateHandleUtil(DataInputStream dis)
             throws IOException {
-        return INSTANCE.deserializeOperatorStateHandle(dis, null);
+        return deserializeOperatorStateHandle(dis, null);
     }
 
     @VisibleForTesting
-    public static void serializeKeyedStateHandleUtil(
-            KeyedStateHandle stateHandle, DataOutputStream dos) throws IOException {
-        INSTANCE.serializeKeyedStateHandle(stateHandle, dos);
+    public void serializeKeyedStateHandleUtil(KeyedStateHandle stateHandle, DataOutputStream dos)
+            throws IOException {
+        serializeKeyedStateHandle(stateHandle, dos);
     }
 
     @VisibleForTesting
-    public static KeyedStateHandle deserializeKeyedStateHandleUtil(DataInputStream dis)
+    public KeyedStateHandle deserializeKeyedStateHandleUtil(DataInputStream dis)
             throws IOException {
-        return INSTANCE.deserializeKeyedStateHandle(dis, null);
+        return deserializeKeyedStateHandle(dis, null);
     }
 
     @VisibleForTesting
-    public static StateObjectCollection<InputStateHandle> deserializeInputStateHandle(
-            DataInputStream dis) throws IOException {
-        return INSTANCE.deserializeInputStateHandle(dis, null);
+    public StateObjectCollection<InputStateHandle> deserializeInputStateHandle(DataInputStream dis)
+            throws IOException {
+        return deserializeInputStateHandle(dis, null);
     }
 
     @VisibleForTesting
     public StateObjectCollection<OutputStateHandle> deserializeOutputStateHandle(
             DataInputStream dis) throws IOException {
-        return INSTANCE.deserializeOutputStateHandle(dis, null);
+        return deserializeOutputStateHandle(dis, null);
     }
 
     static class SubtaskAndFinishedState {
diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/util/OperatorSnapshotUtil.java b/flink-runtime/src/test/java/org/apache/flink/streaming/util/OperatorSnapshotUtil.java
index a2d4ef8beb1..f5eb53923ba 100644
--- a/flink-runtime/src/test/java/org/apache/flink/streaming/util/OperatorSnapshotUtil.java
+++ b/flink-runtime/src/test/java/org/apache/flink/streaming/util/OperatorSnapshotUtil.java
@@ -61,7 +61,8 @@ public class OperatorSnapshotUtil {
             if (rawOperatorState != null) {
                 dos.writeInt(rawOperatorState.size());
                 for (OperatorStateHandle operatorStateHandle : rawOperatorState) {
-                    MetadataV3Serializer.serializeOperatorStateHandleUtil(operatorStateHandle, dos);
+                    MetadataV3Serializer.INSTANCE.serializeOperatorStateHandleUtil(
+                            operatorStateHandle, dos);
                 }
             } else {
                 // this means no states, not even an empty list
@@ -72,7 +73,8 @@ public class OperatorSnapshotUtil {
             if (managedOperatorState != null) {
                 dos.writeInt(managedOperatorState.size());
                 for (OperatorStateHandle operatorStateHandle : managedOperatorState) {
-                    MetadataV3Serializer.serializeOperatorStateHandleUtil(operatorStateHandle, dos);
+                    MetadataV3Serializer.INSTANCE.serializeOperatorStateHandleUtil(
+                            operatorStateHandle, dos);
                 }
             } else {
                 // this means no states, not even an empty list
@@ -83,7 +85,8 @@ public class OperatorSnapshotUtil {
             if (rawKeyedState != null) {
                 dos.writeInt(rawKeyedState.size());
                 for (KeyedStateHandle keyedStateHandle : rawKeyedState) {
-                    MetadataV3Serializer.serializeKeyedStateHandleUtil(keyedStateHandle, dos);
+                    MetadataV3Serializer.INSTANCE.serializeKeyedStateHandleUtil(
+                            keyedStateHandle, dos);
                 }
             } else {
                 // this means no operator states, not even an empty list
@@ -94,7 +97,8 @@ public class OperatorSnapshotUtil {
             if (managedKeyedState != null) {
                 dos.writeInt(managedKeyedState.size());
                 for (KeyedStateHandle keyedStateHandle : managedKeyedState) {
-                    MetadataV3Serializer.serializeKeyedStateHandleUtil(keyedStateHandle, dos);
+                    MetadataV3Serializer.INSTANCE.serializeKeyedStateHandleUtil(
+                            keyedStateHandle, dos);
                 }
             } else {
                 // this means no operator states, not even an empty list
@@ -137,7 +141,7 @@ public class OperatorSnapshotUtil {
                 rawOperatorState = new ArrayList<>();
                 for (int i = 0; i < numRawOperatorStates; i++) {
                     OperatorStateHandle operatorState =
-                            MetadataV3Serializer.deserializeOperatorStateHandleUtil(dis);
+                            MetadataV3Serializer.INSTANCE.deserializeOperatorStateHandleUtil(dis);
                     rawOperatorState.add(operatorState);
                 }
             }
@@ -148,7 +152,7 @@ public class OperatorSnapshotUtil {
                 managedOperatorState = new ArrayList<>();
                 for (int i = 0; i < numManagedOperatorStates; i++) {
                     OperatorStateHandle operatorState =
-                            MetadataV3Serializer.deserializeOperatorStateHandleUtil(dis);
+                            MetadataV3Serializer.INSTANCE.deserializeOperatorStateHandleUtil(dis);
                     managedOperatorState.add(operatorState);
                 }
             }
@@ -159,7 +163,7 @@ public class OperatorSnapshotUtil {
                 rawKeyedState = new ArrayList<>();
                 for (int i = 0; i < numRawKeyedStates; i++) {
                     KeyedStateHandle keyedState =
-                            MetadataV3Serializer.deserializeKeyedStateHandleUtil(dis);
+                            MetadataV3Serializer.INSTANCE.deserializeKeyedStateHandleUtil(dis);
                     rawKeyedState.add(keyedState);
                 }
             }
@@ -170,14 +174,14 @@ public class OperatorSnapshotUtil {
                 managedKeyedState = new ArrayList<>();
                 for (int i = 0; i < numManagedKeyedStates; i++) {
                     KeyedStateHandle keyedState =
-                            MetadataV3Serializer.deserializeKeyedStateHandleUtil(dis);
+                            MetadataV3Serializer.INSTANCE.deserializeKeyedStateHandleUtil(dis);
                     managedKeyedState.add(keyedState);
                 }
             }
 
             final StateObjectCollection<InputStateHandle> inputChannelStateHandles =
                     v == MetadataV3Serializer.VERSION
-                            ? MetadataV3Serializer.deserializeInputStateHandle(dis)
+                            ? MetadataV3Serializer.INSTANCE.deserializeInputStateHandle(dis)
                             : StateObjectCollection.empty();
 
             final StateObjectCollection<OutputStateHandle> resultSubpartitionStateHandles =


(flink) 07/07: [FLINK-33734][checkpointing] Compact and serialize original channel state handle before send checkpoint ack to JobManager

Posted by za...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

zakelly pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit 0d1b5aa1cddb382d852a19c50b67985b7788fc7f
Author: wangfeifan <zo...@163.com>
AuthorDate: Thu Feb 1 18:13:53 2024 +0800

    [FLINK-33734][checkpointing] Compact and serialize original channel state handle before send checkpoint ack to JobManager
---
 .../main/java/org/apache/flink/runtime/checkpoint/Checkpoints.java  | 4 ++--
 .../org/apache/flink/streaming/runtime/tasks/OperatorChain.java     | 6 +++---
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/Checkpoints.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/Checkpoints.java
index 682e9fdcea1..4f1000fd3b3 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/Checkpoints.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/Checkpoints.java
@@ -24,7 +24,7 @@ import org.apache.flink.runtime.OperatorIDPair;
 import org.apache.flink.runtime.checkpoint.metadata.CheckpointMetadata;
 import org.apache.flink.runtime.checkpoint.metadata.MetadataSerializer;
 import org.apache.flink.runtime.checkpoint.metadata.MetadataSerializers;
-import org.apache.flink.runtime.checkpoint.metadata.MetadataV5Serializer;
+import org.apache.flink.runtime.checkpoint.metadata.MetadataV6Serializer;
 import org.apache.flink.runtime.executiongraph.ExecutionJobVertex;
 import org.apache.flink.runtime.jobgraph.JobVertexID;
 import org.apache.flink.runtime.jobgraph.OperatorID;
@@ -85,7 +85,7 @@ public class Checkpoints {
 
     public static void storeCheckpointMetadata(
             CheckpointMetadata checkpointMetadata, DataOutputStream out) throws IOException {
-        storeCheckpointMetadata(checkpointMetadata, out, MetadataV5Serializer.INSTANCE);
+        storeCheckpointMetadata(checkpointMetadata, out, MetadataV6Serializer.INSTANCE);
     }
 
     public static void storeCheckpointMetadata(
diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/OperatorChain.java b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/OperatorChain.java
index febbcb1bdb0..71b2b53586e 100644
--- a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/OperatorChain.java
+++ b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/OperatorChain.java
@@ -27,7 +27,6 @@ import org.apache.flink.metrics.groups.OperatorMetricGroup;
 import org.apache.flink.runtime.checkpoint.CheckpointException;
 import org.apache.flink.runtime.checkpoint.CheckpointMetaData;
 import org.apache.flink.runtime.checkpoint.CheckpointOptions;
-import org.apache.flink.runtime.checkpoint.StateObjectCollection;
 import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter;
 import org.apache.flink.runtime.event.AbstractEvent;
 import org.apache.flink.runtime.execution.Environment;
@@ -44,6 +43,7 @@ import org.apache.flink.runtime.operators.coordination.AcknowledgeCheckpointEven
 import org.apache.flink.runtime.operators.coordination.OperatorEvent;
 import org.apache.flink.runtime.operators.coordination.OperatorEventDispatcher;
 import org.apache.flink.runtime.plugable.SerializationDelegate;
+import org.apache.flink.runtime.state.ChannelStateHelper;
 import org.apache.flink.runtime.state.CheckpointStreamFactory;
 import org.apache.flink.runtime.state.SnapshotResult;
 import org.apache.flink.streaming.api.graph.NonChainedOutput;
@@ -469,14 +469,14 @@ public abstract class OperatorChain<OUT, OP extends StreamOperator<OUT>>
             snapshotInProgress.setInputChannelStateFuture(
                     channelStateWriteResult
                             .getInputChannelStateHandles()
-                            .thenApply(e -> new StateObjectCollection(e))
+                            .thenApply(ChannelStateHelper::mergeInputStateCollection)
                             .thenApply(SnapshotResult::of));
         }
         if (op == getTailOperator()) {
             snapshotInProgress.setResultSubpartitionStateFuture(
                     channelStateWriteResult
                             .getResultSubpartitionStateHandles()
-                            .thenApply(e -> new StateObjectCollection(e))
+                            .thenApply(ChannelStateHelper::mergeOutputStateCollection)
                             .thenApply(SnapshotResult::of));
         }
     }


(flink) 03/07: [FLINK-33734][checkpointing] Extract interface ChannelStateHandleSerializer

Posted by za...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

zakelly pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit db07fe9f67e69ffa4fe982f80c827c6583862075
Author: wangfeifan <zo...@163.com>
AuthorDate: Thu Feb 1 17:04:55 2024 +0800

    [FLINK-33734][checkpointing] Extract interface ChannelStateHandleSerializer
---
 .../metadata/ChannelStateHandleSerializer.java     |  98 ++----------------
 ...er.java => ChannelStateHandleSerializerV1.java} |  14 ++-
 .../metadata/MetadataV2V3SerializerBase.java       |  15 ++-
 .../checkpoint/metadata/MetadataV3Serializer.java  |  40 ++++----
 .../ChannelStateHandleSerializerV1Test.java        | 111 +++++++++++++++++++++
 .../flink/streaming/util/OperatorSnapshotUtil.java |   9 +-
 6 files changed, 160 insertions(+), 127 deletions(-)

diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializer.java
index 150caf1f6f4..3f9dde38f35 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializer.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializer.java
@@ -17,108 +17,28 @@
 
 package org.apache.flink.runtime.checkpoint.metadata;
 
-import org.apache.flink.runtime.state.AbstractChannelStateHandle;
-import org.apache.flink.runtime.state.AbstractChannelStateHandle.StateContentMetaInfo;
-import org.apache.flink.runtime.state.InputChannelStateHandle;
 import org.apache.flink.runtime.state.InputStateHandle;
 import org.apache.flink.runtime.state.OutputStateHandle;
-import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
-import org.apache.flink.runtime.state.StreamStateHandle;
-import org.apache.flink.util.function.BiConsumerWithException;
-import org.apache.flink.util.function.FunctionWithException;
-import org.apache.flink.util.function.QuadFunction;
 
 import java.io.DataInputStream;
 import java.io.DataOutputStream;
 import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
-
-import static org.apache.flink.runtime.checkpoint.metadata.MetadataV2V3SerializerBase.deserializeStreamStateHandle;
-import static org.apache.flink.runtime.checkpoint.metadata.MetadataV2V3SerializerBase.serializeStreamStateHandle;
-import static org.apache.flink.runtime.state.ChannelStateHelper.INPUT_CHANNEL_INFO_READER;
-import static org.apache.flink.runtime.state.ChannelStateHelper.INPUT_CHANNEL_INFO_WRITER;
-import static org.apache.flink.runtime.state.ChannelStateHelper.RESULT_SUBPARTITION_INFO_READER;
-import static org.apache.flink.runtime.state.ChannelStateHelper.RESULT_SUBPARTITION_INFO_WRITER;
 
 /**
- * ChannelStateHandlerSerializer only support {@link InputChannelStateHandle} and {@link
- * ResultSubpartitionStateHandle}, used in {@link MetadataV3Serializer} and {@link
- * MetadataV4Serializer}.
+ * (De)serializer for unaligned checkpoint channel state handle which is used in {@link
+ * MetadataSerializer} .
  */
-class ChannelStateHandleSerializer {
-
-    public void serialize(OutputStateHandle handle, DataOutputStream dos) throws IOException {
-        if (!(handle instanceof ResultSubpartitionStateHandle)) {
-            throw new IllegalStateException(
-                    "OutputStateHandle must be ResultSubpartitionStateHandle.");
-        }
+interface ChannelStateHandleSerializer {
 
-        serializeChannelStateHandle(
-                (ResultSubpartitionStateHandle) handle, dos, RESULT_SUBPARTITION_INFO_WRITER);
-    }
+    void serialize(OutputStateHandle handle, DataOutputStream dos) throws IOException;
 
-    public OutputStateHandle deserializeOutputStateHandle(
+    OutputStateHandle deserializeOutputStateHandle(
             DataInputStream dis, MetadataV2V3SerializerBase.DeserializationContext context)
-            throws IOException {
-
-        return deserializeChannelStateHandle(
-                RESULT_SUBPARTITION_INFO_READER, ResultSubpartitionStateHandle::new, dis, context);
-    }
-
-    public void serialize(InputStateHandle handle, DataOutputStream dos) throws IOException {
-        if (!(handle instanceof InputChannelStateHandle)) {
-            throw new IllegalStateException("InputStateHandle must be InputChannelStateHandle.");
-        }
+            throws IOException;
 
-        serializeChannelStateHandle(
-                (InputChannelStateHandle) handle, dos, INPUT_CHANNEL_INFO_WRITER);
-    }
+    void serialize(InputStateHandle handle, DataOutputStream dos) throws IOException;
 
-    public InputStateHandle deserializeInputStateHandle(
+    InputStateHandle deserializeInputStateHandle(
             DataInputStream dis, MetadataV2V3SerializerBase.DeserializationContext context)
-            throws IOException {
-
-        return deserializeChannelStateHandle(
-                INPUT_CHANNEL_INFO_READER, InputChannelStateHandle::new, dis, context);
-    }
-
-    private static <I> void serializeChannelStateHandle(
-            AbstractChannelStateHandle<I> handle,
-            DataOutputStream dos,
-            BiConsumerWithException<I, DataOutputStream, IOException> infoWriter)
-            throws IOException {
-        dos.writeInt(handle.getSubtaskIndex());
-        infoWriter.accept(handle.getInfo(), dos);
-        dos.writeInt(handle.getOffsets().size());
-        for (long offset : handle.getOffsets()) {
-            dos.writeLong(offset);
-        }
-        dos.writeLong(handle.getStateSize());
-        serializeStreamStateHandle(handle.getDelegate(), dos);
-    }
-
-    private static <Info, Handle extends AbstractChannelStateHandle<Info>>
-            Handle deserializeChannelStateHandle(
-                    FunctionWithException<DataInputStream, Info, IOException> infoReader,
-                    QuadFunction<Integer, Info, StreamStateHandle, StateContentMetaInfo, Handle>
-                            handleBuilder,
-                    DataInputStream dis,
-                    MetadataV2V3SerializerBase.DeserializationContext context)
-                    throws IOException {
-
-        int subtaskIndex = dis.readInt();
-        final Info info = infoReader.apply(dis);
-        int offsetsSize = dis.readInt();
-        final List<Long> offsets = new ArrayList<>(offsetsSize);
-        for (int i = 0; i < offsetsSize; i++) {
-            offsets.add(dis.readLong());
-        }
-        final long size = dis.readLong();
-        return handleBuilder.apply(
-                subtaskIndex,
-                info,
-                deserializeStreamStateHandle(dis, context),
-                new StateContentMetaInfo(offsets, size));
-    }
+            throws IOException;
 }
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV1.java
similarity index 93%
copy from flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializer.java
copy to flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV1.java
index 150caf1f6f4..22eefc63cf0 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializer.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV1.java
@@ -46,8 +46,9 @@ import static org.apache.flink.runtime.state.ChannelStateHelper.RESULT_SUBPARTIT
  * ResultSubpartitionStateHandle}, used in {@link MetadataV3Serializer} and {@link
  * MetadataV4Serializer}.
  */
-class ChannelStateHandleSerializer {
+class ChannelStateHandleSerializerV1 implements ChannelStateHandleSerializer {
 
+    @Override
     public void serialize(OutputStateHandle handle, DataOutputStream dos) throws IOException {
         if (!(handle instanceof ResultSubpartitionStateHandle)) {
             throw new IllegalStateException(
@@ -58,6 +59,7 @@ class ChannelStateHandleSerializer {
                 (ResultSubpartitionStateHandle) handle, dos, RESULT_SUBPARTITION_INFO_WRITER);
     }
 
+    @Override
     public OutputStateHandle deserializeOutputStateHandle(
             DataInputStream dis, MetadataV2V3SerializerBase.DeserializationContext context)
             throws IOException {
@@ -66,6 +68,7 @@ class ChannelStateHandleSerializer {
                 RESULT_SUBPARTITION_INFO_READER, ResultSubpartitionStateHandle::new, dis, context);
     }
 
+    @Override
     public void serialize(InputStateHandle handle, DataOutputStream dos) throws IOException {
         if (!(handle instanceof InputChannelStateHandle)) {
             throw new IllegalStateException("InputStateHandle must be InputChannelStateHandle.");
@@ -75,6 +78,7 @@ class ChannelStateHandleSerializer {
                 (InputChannelStateHandle) handle, dos, INPUT_CHANNEL_INFO_WRITER);
     }
 
+    @Override
     public InputStateHandle deserializeInputStateHandle(
             DataInputStream dis, MetadataV2V3SerializerBase.DeserializationContext context)
             throws IOException {
@@ -83,10 +87,10 @@ class ChannelStateHandleSerializer {
                 INPUT_CHANNEL_INFO_READER, InputChannelStateHandle::new, dis, context);
     }
 
-    private static <I> void serializeChannelStateHandle(
-            AbstractChannelStateHandle<I> handle,
+    static <Info> void serializeChannelStateHandle(
+            AbstractChannelStateHandle<Info> handle,
             DataOutputStream dos,
-            BiConsumerWithException<I, DataOutputStream, IOException> infoWriter)
+            BiConsumerWithException<Info, DataOutputStream, IOException> infoWriter)
             throws IOException {
         dos.writeInt(handle.getSubtaskIndex());
         infoWriter.accept(handle.getInfo(), dos);
@@ -98,7 +102,7 @@ class ChannelStateHandleSerializer {
         serializeStreamStateHandle(handle.getDelegate(), dos);
     }
 
-    private static <Info, Handle extends AbstractChannelStateHandle<Info>>
+    static <Info, Handle extends AbstractChannelStateHandle<Info>>
             Handle deserializeChannelStateHandle(
                     FunctionWithException<DataInputStream, Info, IOException> infoReader,
                     QuadFunction<Integer, Info, StreamStateHandle, StateContentMetaInfo, Handle>
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV2V3SerializerBase.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV2V3SerializerBase.java
index 37d043047d5..57d69b29f3e 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV2V3SerializerBase.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV2V3SerializerBase.java
@@ -312,8 +312,8 @@ public abstract class MetadataV2V3SerializerBase {
             state.setRawKeyedState(rawKeyedState);
         }
 
-        state.setInputChannelState(deserializeInputChannelStateHandle(dis, context));
-        state.setResultSubpartitionState(deserializeResultSubpartitionStateHandle(dis, context));
+        state.setInputChannelState(deserializeInputStateHandle(dis, context));
+        state.setResultSubpartitionState(deserializeOutputStateHandle(dis, context));
 
         return state.build();
     }
@@ -696,22 +696,21 @@ public abstract class MetadataV2V3SerializerBase {
     //  channel state (unaligned checkpoints)
     // ------------------------------------------------------------------------
 
-    protected StateObjectCollection<OutputStateHandle> deserializeResultSubpartitionStateHandle(
+    protected StateObjectCollection<OutputStateHandle> deserializeOutputStateHandle(
             DataInputStream dis, @Nullable DeserializationContext context) throws IOException {
         return StateObjectCollection.empty();
     }
 
-    protected StateObjectCollection<InputStateHandle> deserializeInputChannelStateHandle(
+    protected StateObjectCollection<InputStateHandle> deserializeInputStateHandle(
             DataInputStream dis, @Nullable DeserializationContext context) throws IOException {
         return StateObjectCollection.empty();
     }
 
-    protected void serializeResultSubpartitionStateHandle(
-            OutputStateHandle resultSubpartitionStateHandle, DataOutputStream dos)
+    protected void serializeOutputStateHandle(OutputStateHandle handle, DataOutputStream dos)
             throws IOException {}
 
-    protected void serializeInputChannelStateHandle(
-            InputStateHandle inputChannelStateHandle, DataOutputStream dos) throws IOException {}
+    protected void serializeInputStateHandle(InputStateHandle handle, DataOutputStream dos)
+            throws IOException {}
 
     // ------------------------------------------------------------------------
     //  low-level state handles
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV3Serializer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV3Serializer.java
index c2fe5134e08..ed257b9d64d 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV3Serializer.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV3Serializer.java
@@ -63,7 +63,7 @@ public class MetadataV3Serializer extends MetadataV2V3SerializerBase implements
     public static final MetadataV3Serializer INSTANCE = new MetadataV3Serializer();
 
     private final ChannelStateHandleSerializer channelStateHandleSerializer =
-            new ChannelStateHandleSerializer();
+            new ChannelStateHandleSerializerV1();
 
     /** Singleton, not meant to be instantiated. */
     protected MetadataV3Serializer() {}
@@ -141,11 +141,9 @@ public class MetadataV3Serializer extends MetadataV2V3SerializerBase implements
             throws IOException {
         super.serializeSubtaskState(subtaskState, dos);
         serializeCollection(
-                subtaskState.getInputChannelState(), dos, this::serializeInputChannelStateHandle);
+                subtaskState.getInputChannelState(), dos, this::serializeInputStateHandle);
         serializeCollection(
-                subtaskState.getResultSubpartitionState(),
-                dos,
-                this::serializeResultSubpartitionStateHandle);
+                subtaskState.getResultSubpartitionState(), dos, this::serializeOutputStateHandle);
     }
 
     @Override
@@ -200,34 +198,36 @@ public class MetadataV3Serializer extends MetadataV2V3SerializerBase implements
         }
     }
 
+    protected ChannelStateHandleSerializer getChannelStateHandleSerializer() {
+        return channelStateHandleSerializer;
+    }
+
     @VisibleForTesting
     @Override
-    public void serializeResultSubpartitionStateHandle(
-            OutputStateHandle handle, DataOutputStream dos) throws IOException {
-        channelStateHandleSerializer.serialize(handle, dos);
+    public void serializeOutputStateHandle(OutputStateHandle handle, DataOutputStream dos)
+            throws IOException {
+        getChannelStateHandleSerializer().serialize(handle, dos);
     }
 
-    @VisibleForTesting
     @Override
-    public StateObjectCollection<OutputStateHandle> deserializeResultSubpartitionStateHandle(
+    public StateObjectCollection<OutputStateHandle> deserializeOutputStateHandle(
             DataInputStream dis, @Nullable DeserializationContext context) throws IOException {
         return deserializeCollection(
-                dis, context, channelStateHandleSerializer::deserializeOutputStateHandle);
+                dis, context, getChannelStateHandleSerializer()::deserializeOutputStateHandle);
     }
 
     @VisibleForTesting
     @Override
-    public void serializeInputChannelStateHandle(InputStateHandle handle, DataOutputStream dos)
+    public void serializeInputStateHandle(InputStateHandle handle, DataOutputStream dos)
             throws IOException {
-        channelStateHandleSerializer.serialize(handle, dos);
+        getChannelStateHandleSerializer().serialize(handle, dos);
     }
 
-    @VisibleForTesting
     @Override
-    public StateObjectCollection<InputStateHandle> deserializeInputChannelStateHandle(
+    public StateObjectCollection<InputStateHandle> deserializeInputStateHandle(
             DataInputStream dis, @Nullable DeserializationContext context) throws IOException {
         return deserializeCollection(
-                dis, context, channelStateHandleSerializer::deserializeInputStateHandle);
+                dis, context, getChannelStateHandleSerializer()::deserializeInputStateHandle);
     }
 
     private <T extends StateObject> void serializeCollection(
@@ -292,15 +292,15 @@ public class MetadataV3Serializer extends MetadataV2V3SerializerBase implements
     }
 
     @VisibleForTesting
-    public static StateObjectCollection<InputStateHandle> deserializeInputChannelStateHandle(
+    public static StateObjectCollection<InputStateHandle> deserializeInputStateHandle(
             DataInputStream dis) throws IOException {
-        return INSTANCE.deserializeInputChannelStateHandle(dis, null);
+        return INSTANCE.deserializeInputStateHandle(dis, null);
     }
 
     @VisibleForTesting
-    public StateObjectCollection<OutputStateHandle> deserializeResultSubpartitionStateHandle(
+    public StateObjectCollection<OutputStateHandle> deserializeOutputStateHandle(
             DataInputStream dis) throws IOException {
-        return INSTANCE.deserializeResultSubpartitionStateHandle(dis, null);
+        return INSTANCE.deserializeOutputStateHandle(dis, null);
     }
 
     static class SubtaskAndFinishedState {
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV1Test.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV1Test.java
new file mode 100644
index 00000000000..5823cf5beaa
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV1Test.java
@@ -0,0 +1,111 @@
+/*
+ * 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.metadata;
+
+import org.apache.flink.api.java.tuple.Tuple5;
+import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo;
+import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo;
+import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ThreadLocalRandom;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+class ChannelStateHandleSerializerV1Test {
+
+    @Test
+    void testSerializeInputHandle() throws IOException {
+        ChannelStateHandleSerializerV1 serializer = new ChannelStateHandleSerializerV1();
+
+        Tuple5<Integer, Integer, Integer, StreamStateHandle, List<Long>> handleFields =
+                generateRandomChannelHandleFields();
+
+        InputChannelStateHandle inputHandle =
+                new InputChannelStateHandle(
+                        handleFields.f0,
+                        new InputChannelInfo(handleFields.f1, handleFields.f2),
+                        handleFields.f3,
+                        handleFields.f4,
+                        handleFields.f3.getStateSize());
+
+        ByteArrayOutputStream inMemBuffer = new ByteArrayOutputStream();
+        DataOutputStream outputStream = new DataOutputStream(inMemBuffer);
+        serializer.serialize(inputHandle, outputStream);
+
+        DataInputStream inputStream =
+                new DataInputStream(new ByteArrayInputStream(inMemBuffer.toByteArray()));
+        assertEquals(inputHandle, serializer.deserializeInputStateHandle(inputStream, null));
+    }
+
+    @Test
+    void testSerializeOutputHandle() throws IOException {
+        ChannelStateHandleSerializerV1 serializer = new ChannelStateHandleSerializerV1();
+
+        Tuple5<Integer, Integer, Integer, StreamStateHandle, List<Long>> handleFields =
+                generateRandomChannelHandleFields();
+
+        ResultSubpartitionStateHandle outputHandle =
+                new ResultSubpartitionStateHandle(
+                        handleFields.f0,
+                        new ResultSubpartitionInfo(handleFields.f1, handleFields.f2),
+                        handleFields.f3,
+                        handleFields.f4,
+                        handleFields.f3.getStateSize());
+
+        ByteArrayOutputStream inMemBuffer = new ByteArrayOutputStream();
+        DataOutputStream outputStream = new DataOutputStream(inMemBuffer);
+        serializer.serialize(outputHandle, outputStream);
+
+        DataInputStream inputStream =
+                new DataInputStream(new ByteArrayInputStream(inMemBuffer.toByteArray()));
+        assertEquals(outputHandle, serializer.deserializeInputStateHandle(inputStream, null));
+    }
+
+    private Tuple5<Integer, Integer, Integer, StreamStateHandle, List<Long>>
+            generateRandomChannelHandleFields() {
+        ThreadLocalRandom random = ThreadLocalRandom.current();
+
+        int subtaskIndex = random.nextInt();
+        int infoIdx1 = random.nextInt();
+        int infoIdx2 = random.nextInt();
+
+        byte[] randomBytes = String.valueOf(random.nextLong()).getBytes();
+        StreamStateHandle delegate = new ByteStreamStateHandle("delegated", randomBytes);
+
+        int offsetSize = random.nextInt(10, 20);
+        List<Long> offsets = new ArrayList<>(offsetSize);
+        for (int i = 0; i < offsetSize; i++) {
+            offsets.add(random.nextLong());
+        }
+
+        return Tuple5.of(subtaskIndex, infoIdx1, infoIdx2, delegate, offsets);
+    }
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/util/OperatorSnapshotUtil.java b/flink-runtime/src/test/java/org/apache/flink/streaming/util/OperatorSnapshotUtil.java
index 7f9312f0c38..a2d4ef8beb1 100644
--- a/flink-runtime/src/test/java/org/apache/flink/streaming/util/OperatorSnapshotUtil.java
+++ b/flink-runtime/src/test/java/org/apache/flink/streaming/util/OperatorSnapshotUtil.java
@@ -104,7 +104,7 @@ public class OperatorSnapshotUtil {
             Collection<InputStateHandle> inputChannelStateHandles = state.getInputChannelState();
             dos.writeInt(inputChannelStateHandles.size());
             for (InputStateHandle inputChannelStateHandle : inputChannelStateHandles) {
-                MetadataV3Serializer.INSTANCE.serializeInputChannelStateHandle(
+                MetadataV3Serializer.INSTANCE.serializeInputStateHandle(
                         inputChannelStateHandle, dos);
             }
 
@@ -112,7 +112,7 @@ public class OperatorSnapshotUtil {
                     state.getResultSubpartitionState();
             dos.writeInt(inputChannelStateHandles.size());
             for (OutputStateHandle resultSubpartitionStateHandle : resultSubpartitionStateHandles) {
-                MetadataV3Serializer.INSTANCE.serializeResultSubpartitionStateHandle(
+                MetadataV3Serializer.INSTANCE.serializeOutputStateHandle(
                         resultSubpartitionStateHandle, dos);
             }
 
@@ -177,13 +177,12 @@ public class OperatorSnapshotUtil {
 
             final StateObjectCollection<InputStateHandle> inputChannelStateHandles =
                     v == MetadataV3Serializer.VERSION
-                            ? MetadataV3Serializer.deserializeInputChannelStateHandle(dis)
+                            ? MetadataV3Serializer.deserializeInputStateHandle(dis)
                             : StateObjectCollection.empty();
 
             final StateObjectCollection<OutputStateHandle> resultSubpartitionStateHandles =
                     v == MetadataV3Serializer.VERSION
-                            ? MetadataV3Serializer.INSTANCE
-                                    .deserializeResultSubpartitionStateHandle(dis)
+                            ? MetadataV3Serializer.INSTANCE.deserializeOutputStateHandle(dis)
                             : StateObjectCollection.empty();
 
             return OperatorSubtaskState.builder()


(flink) 01/07: [FLINK-33734][checkpointing] Introduce interface ChannelState,InputStateHandle,OutputStateHandle

Posted by za...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

zakelly pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit be91c9f278f3543089facfd7b0a97188d30a65f1
Author: wangfeifan <zo...@163.com>
AuthorDate: Thu Feb 1 12:23:43 2024 +0800

    [FLINK-33734][checkpointing] Introduce interface ChannelState,InputStateHandle,OutputStateHandle
---
 .../apache/flink/runtime/state/ChannelState.java   | 22 ++++++++++++++++++++++
 .../runtime/state/InputChannelStateHandle.java     |  3 ++-
 .../flink/runtime/state/InputStateHandle.java      | 22 ++++++++++++++++++++++
 .../flink/runtime/state/OutputStateHandle.java     | 22 ++++++++++++++++++++++
 .../state/ResultSubpartitionStateHandle.java       |  2 +-
 5 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/ChannelState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/ChannelState.java
new file mode 100644
index 00000000000..33e1ab7b6b7
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/ChannelState.java
@@ -0,0 +1,22 @@
+/*
+ * 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.state;
+
+/** empty interface just for mark a state object is unaligned checkpoint inflight-data. */
+public interface ChannelState extends StateObject {}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/InputChannelStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/InputChannelStateHandle.java
index 6fdf8b84999..36f36ebb61c 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/InputChannelStateHandle.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/InputChannelStateHandle.java
@@ -28,7 +28,8 @@ import java.util.List;
  * org.apache.flink.runtime.io.network.partition.consumer.InputChannel InputChannel} state.
  */
 @Internal
-public class InputChannelStateHandle extends AbstractChannelStateHandle<InputChannelInfo> {
+public class InputChannelStateHandle extends AbstractChannelStateHandle<InputChannelInfo>
+        implements InputStateHandle {
 
     private static final long serialVersionUID = 1L;
 
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/InputStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/InputStateHandle.java
new file mode 100644
index 00000000000..eb7b4c7de66
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/InputStateHandle.java
@@ -0,0 +1,22 @@
+/*
+ * 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.state;
+
+/** empty interface just for mark a state object is input inflight-data. */
+public interface InputStateHandle extends ChannelState {}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/OutputStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/OutputStateHandle.java
new file mode 100644
index 00000000000..e6c5f205d29
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/OutputStateHandle.java
@@ -0,0 +1,22 @@
+/*
+ * 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.state;
+
+/** empty interface just for mark a state object is output inflight-data. */
+public interface OutputStateHandle extends ChannelState {}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/ResultSubpartitionStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/ResultSubpartitionStateHandle.java
index 0f281506f49..4b969e028ed 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/ResultSubpartitionStateHandle.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/ResultSubpartitionStateHandle.java
@@ -29,7 +29,7 @@ import java.util.List;
  */
 @Internal
 public class ResultSubpartitionStateHandle
-        extends AbstractChannelStateHandle<ResultSubpartitionInfo> {
+        extends AbstractChannelStateHandle<ResultSubpartitionInfo> implements OutputStateHandle {
 
     private static final long serialVersionUID = 1L;
 


(flink) 02/07: [FLINK-33734][checkpointing] Use interface replace InputChannelStateHandle and ResultSubpartitionStateHandle in OperatorSubtaskState

Posted by za...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

zakelly pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit 32aa7973daaf70e097f247eaf5a3869b47d8e3e0
Author: wangfeifan <zo...@163.com>
AuthorDate: Thu Feb 1 16:10:33 2024 +0800

    [FLINK-33734][checkpointing] Use interface replace InputChannelStateHandle and ResultSubpartitionStateHandle in OperatorSubtaskState
---
 .../runtime/checkpoint/OperatorSubtaskState.java   |  37 +++---
 .../PrioritizedOperatorSubtaskState.java           |  94 +++++++++++---
 .../checkpoint/StateAssignmentOperation.java       |   9 +-
 .../runtime/checkpoint/TaskStateAssignment.java    |  12 +-
 .../channel/SequentialChannelStateReaderImpl.java  |   6 +-
 .../metadata/ChannelStateHandleSerializer.java     |  55 ++++----
 .../metadata/MetadataV2V3SerializerBase.java       |  17 ++-
 .../checkpoint/metadata/MetadataV3Serializer.java  |  30 ++---
 .../flink/runtime/state/ChannelStateHelper.java    | 140 +++++++++++++++++++++
 .../api/operators/OperatorSnapshotFinalizer.java   |   8 +-
 .../api/operators/OperatorSnapshotFutures.java     |  21 ++--
 .../streaming/runtime/tasks/OperatorChain.java     |   4 +-
 .../CheckpointCoordinatorFailureTest.java          |  15 ++-
 .../SequentialChannelStateReaderImplTest.java      |   7 +-
 .../runtime/state/ChannelPersistenceITCase.java    |   7 +-
 .../operators/OperatorSnapshotFinalizerTest.java   |   6 +-
 .../api/operators/OperatorSnapshotFuturesTest.java |  17 +--
 .../operators/StreamOperatorStateHandlerTest.java  |   8 +-
 .../flink/streaming/util/OperatorSnapshotUtil.java |  27 ++--
 .../runtime/tasks/LocalStateForwardingTest.java    |   9 +-
 20 files changed, 367 insertions(+), 162 deletions(-)

diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/OperatorSubtaskState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/OperatorSubtaskState.java
index 24ff9db5f04..a70b85175d6 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/OperatorSubtaskState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/OperatorSubtaskState.java
@@ -19,12 +19,12 @@
 package org.apache.flink.runtime.checkpoint;
 
 import org.apache.flink.annotation.VisibleForTesting;
-import org.apache.flink.runtime.state.AbstractChannelStateHandle;
+import org.apache.flink.runtime.state.ChannelState;
 import org.apache.flink.runtime.state.CompositeStateHandle;
-import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.InputStateHandle;
 import org.apache.flink.runtime.state.KeyedStateHandle;
 import org.apache.flink.runtime.state.OperatorStateHandle;
-import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
+import org.apache.flink.runtime.state.OutputStateHandle;
 import org.apache.flink.runtime.state.SharedStateRegistry;
 import org.apache.flink.runtime.state.SharedStateRegistryImpl.EmptyDiscardStateObjectForRegister;
 import org.apache.flink.runtime.state.SharedStateRegistryKey;
@@ -42,7 +42,7 @@ import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
-import static org.apache.flink.runtime.state.AbstractChannelStateHandle.collectUniqueDelegates;
+import static org.apache.flink.runtime.state.ChannelStateHelper.collectUniqueDisposableInChannelState;
 import static org.apache.flink.util.Preconditions.checkNotNull;
 
 /**
@@ -87,9 +87,9 @@ public class OperatorSubtaskState implements CompositeStateHandle {
      */
     private final StateObjectCollection<KeyedStateHandle> rawKeyedState;
 
-    private final StateObjectCollection<InputChannelStateHandle> inputChannelState;
+    private final StateObjectCollection<InputStateHandle> inputChannelState;
 
-    private final StateObjectCollection<ResultSubpartitionStateHandle> resultSubpartitionState;
+    private final StateObjectCollection<OutputStateHandle> resultSubpartitionState;
 
     /**
      * The subpartitions mappings per partition set when the output operator for a partition was
@@ -122,8 +122,8 @@ public class OperatorSubtaskState implements CompositeStateHandle {
             StateObjectCollection<OperatorStateHandle> rawOperatorState,
             StateObjectCollection<KeyedStateHandle> managedKeyedState,
             StateObjectCollection<KeyedStateHandle> rawKeyedState,
-            StateObjectCollection<InputChannelStateHandle> inputChannelState,
-            StateObjectCollection<ResultSubpartitionStateHandle> resultSubpartitionState,
+            StateObjectCollection<InputStateHandle> inputChannelState,
+            StateObjectCollection<OutputStateHandle> resultSubpartitionState,
             InflightDataRescalingDescriptor inputRescalingDescriptor,
             InflightDataRescalingDescriptor outputRescalingDescriptor) {
 
@@ -151,11 +151,8 @@ public class OperatorSubtaskState implements CompositeStateHandle {
                 .filter(Objects::nonNull);
     }
 
-    private Stream<StateObjectCollection<? extends AbstractChannelStateHandle<?>>>
-            streamChannelStates() {
-        return Stream.<StateObjectCollection<? extends AbstractChannelStateHandle<?>>>of(
-                        inputChannelState, resultSubpartitionState)
-                .filter(Objects::nonNull);
+    private Stream<StateObjectCollection<? extends ChannelState>> streamChannelStates() {
+        return Stream.of(inputChannelState, resultSubpartitionState).filter(Objects::nonNull);
     }
 
     @VisibleForTesting
@@ -189,11 +186,11 @@ public class OperatorSubtaskState implements CompositeStateHandle {
         return rawKeyedState;
     }
 
-    public StateObjectCollection<InputChannelStateHandle> getInputChannelState() {
+    public StateObjectCollection<InputStateHandle> getInputChannelState() {
         return inputChannelState;
     }
 
-    public StateObjectCollection<ResultSubpartitionStateHandle> getResultSubpartitionState() {
+    public StateObjectCollection<OutputStateHandle> getResultSubpartitionState() {
         return resultSubpartitionState;
     }
 
@@ -208,7 +205,7 @@ public class OperatorSubtaskState implements CompositeStateHandle {
     public List<StateObject> getDiscardables() {
         return Stream.concat(
                         streamOperatorAndKeyedStates().flatMap(Collection::stream),
-                        collectUniqueDelegates(streamChannelStates()))
+                        collectUniqueDisposableInChannelState(streamChannelStates()))
                 .collect(Collectors.toList());
     }
 
@@ -393,9 +390,9 @@ public class OperatorSubtaskState implements CompositeStateHandle {
                 StateObjectCollection.empty();
         private StateObjectCollection<KeyedStateHandle> rawKeyedState =
                 StateObjectCollection.empty();
-        private StateObjectCollection<InputChannelStateHandle> inputChannelState =
+        private StateObjectCollection<InputStateHandle> inputChannelState =
                 StateObjectCollection.empty();
-        private StateObjectCollection<ResultSubpartitionStateHandle> resultSubpartitionState =
+        private StateObjectCollection<OutputStateHandle> resultSubpartitionState =
                 StateObjectCollection.empty();
         private InflightDataRescalingDescriptor inputRescalingDescriptor =
                 InflightDataRescalingDescriptor.NO_RESCALE;
@@ -447,13 +444,13 @@ public class OperatorSubtaskState implements CompositeStateHandle {
         }
 
         public Builder setInputChannelState(
-                StateObjectCollection<InputChannelStateHandle> inputChannelState) {
+                StateObjectCollection<InputStateHandle> inputChannelState) {
             this.inputChannelState = checkNotNull(inputChannelState);
             return this;
         }
 
         public Builder setResultSubpartitionState(
-                StateObjectCollection<ResultSubpartitionStateHandle> resultSubpartitionState) {
+                StateObjectCollection<OutputStateHandle> resultSubpartitionState) {
             this.resultSubpartitionState = checkNotNull(resultSubpartitionState);
             return this;
         }
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PrioritizedOperatorSubtaskState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PrioritizedOperatorSubtaskState.java
index e41bcfe7338..3e1aa039f0d 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PrioritizedOperatorSubtaskState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PrioritizedOperatorSubtaskState.java
@@ -19,10 +19,14 @@
 package org.apache.flink.runtime.checkpoint;
 
 import org.apache.flink.annotation.Internal;
-import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo;
+import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo;
+import org.apache.flink.runtime.state.AbstractChannelStateHandle;
+import org.apache.flink.runtime.state.ChannelState;
+import org.apache.flink.runtime.state.InputStateHandle;
 import org.apache.flink.runtime.state.KeyedStateHandle;
 import org.apache.flink.runtime.state.OperatorStateHandle;
-import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
+import org.apache.flink.runtime.state.OutputStateHandle;
 import org.apache.flink.runtime.state.StateObject;
 
 import org.apache.commons.lang3.BooleanUtils;
@@ -33,6 +37,8 @@ import javax.annotation.Nullable;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Comparator;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -73,10 +79,9 @@ public class PrioritizedOperatorSubtaskState {
     /** List of prioritized snapshot alternatives for raw keyed state. */
     private final List<StateObjectCollection<KeyedStateHandle>> prioritizedRawKeyedState;
 
-    private final List<StateObjectCollection<InputChannelStateHandle>> prioritizedInputChannelState;
+    private final List<StateObjectCollection<InputStateHandle>> prioritizedInputChannelState;
 
-    private final List<StateObjectCollection<ResultSubpartitionStateHandle>>
-            prioritizedResultSubpartitionState;
+    private final List<StateObjectCollection<OutputStateHandle>> prioritizedResultSubpartitionState;
 
     /** Checkpoint id for a restored operator or null if not restored. */
     private final @Nullable Long restoredCheckpointId;
@@ -88,11 +93,9 @@ public class PrioritizedOperatorSubtaskState {
                     List<StateObjectCollection<OperatorStateHandle>>
                             prioritizedManagedOperatorState,
             @Nonnull List<StateObjectCollection<OperatorStateHandle>> prioritizedRawOperatorState,
+            @Nonnull List<StateObjectCollection<InputStateHandle>> prioritizedInputChannelState,
             @Nonnull
-                    List<StateObjectCollection<InputChannelStateHandle>>
-                            prioritizedInputChannelState,
-            @Nonnull
-                    List<StateObjectCollection<ResultSubpartitionStateHandle>>
+                    List<StateObjectCollection<OutputStateHandle>>
                             prioritizedResultSubpartitionState,
             @Nullable Long restoredCheckpointId) {
 
@@ -182,13 +185,12 @@ public class PrioritizedOperatorSubtaskState {
     }
 
     @Nonnull
-    public StateObjectCollection<InputChannelStateHandle> getPrioritizedInputChannelState() {
+    public StateObjectCollection<InputStateHandle> getPrioritizedInputChannelState() {
         return lastElement(prioritizedInputChannelState);
     }
 
     @Nonnull
-    public StateObjectCollection<ResultSubpartitionStateHandle>
-            getPrioritizedResultSubpartitionState() {
+    public StateObjectCollection<OutputStateHandle> getPrioritizedResultSubpartitionState() {
         return lastElement(prioritizedResultSubpartitionState);
     }
 
@@ -271,10 +273,10 @@ public class PrioritizedOperatorSubtaskState {
                     new ArrayList<>(size);
             List<StateObjectCollection<KeyedStateHandle>> rawKeyedAlternatives =
                     new ArrayList<>(size);
-            List<StateObjectCollection<InputChannelStateHandle>> inputChannelStateAlternatives =
+            List<StateObjectCollection<InputStateHandle>> inputChannelStateAlternatives =
+                    new ArrayList<>(size);
+            List<StateObjectCollection<OutputStateHandle>> resultSubpartitionStateAlternatives =
                     new ArrayList<>(size);
-            List<StateObjectCollection<ResultSubpartitionStateHandle>>
-                    resultSubpartitionStateAlternatives = new ArrayList<>(size);
 
             for (OperatorSubtaskState subtaskState : alternativesByPriority) {
 
@@ -309,11 +311,31 @@ public class PrioritizedOperatorSubtaskState {
                     resolvePrioritizedAlternatives(
                             jobManagerState.getInputChannelState(),
                             inputChannelStateAlternatives,
-                            eqStateApprover(InputChannelStateHandle::getInfo)),
+                            channelStateApprover(
+                                    (InputChannelInfo i1, InputChannelInfo i2) -> {
+                                        if (i1.getGateIdx() == i2.getGateIdx()) {
+                                            return Integer.compare(
+                                                    i1.getInputChannelIdx(),
+                                                    i2.getInputChannelIdx());
+                                        } else {
+                                            return Integer.compare(
+                                                    i1.getGateIdx(), i2.getGateIdx());
+                                        }
+                                    })),
                     resolvePrioritizedAlternatives(
                             jobManagerState.getResultSubpartitionState(),
                             resultSubpartitionStateAlternatives,
-                            eqStateApprover(ResultSubpartitionStateHandle::getInfo)),
+                            channelStateApprover(
+                                    (ResultSubpartitionInfo r1, ResultSubpartitionInfo r2) -> {
+                                        if (r1.getPartitionIdx() == r2.getPartitionIdx()) {
+                                            return Integer.compare(
+                                                    r1.getSubPartitionIdx(),
+                                                    r2.getSubPartitionIdx());
+                                        } else {
+                                            return Integer.compare(
+                                                    r1.getPartitionIdx(), r2.getPartitionIdx());
+                                        }
+                                    })),
                     restoredCheckpointId);
         }
 
@@ -467,6 +489,44 @@ public class PrioritizedOperatorSubtaskState {
         }
     }
 
+    private static <T extends ChannelState, Info> BiFunction<T, T, Boolean> channelStateApprover(
+            Comparator<Info> comparator) {
+        return (ref, alt) -> approveChannelState(ref, alt, comparator);
+    }
+
+    private static <Info> boolean approveChannelState(
+            ChannelState ref, ChannelState alt, Comparator<Info> comparator) {
+        List<Info> refInfos = extractInfos(ref, comparator);
+        List<Info> altInfos = extractInfos(alt, comparator);
+
+        if (refInfos.size() != altInfos.size()) {
+            return false;
+        }
+
+        Iterator<Info> refInfoIterator = refInfos.iterator();
+        Iterator<Info> altInfoIterator = altInfos.iterator();
+        while (refInfoIterator.hasNext()) {
+            Info refInfo = refInfoIterator.next();
+            Info altInfo = altInfoIterator.next();
+            if (!refInfo.equals(altInfo)) {
+                return false;
+            }
+        }
+
+        return true;
+    }
+
+    /** return sorted info list extracted from the passed in handle. */
+    private static <Info> List<Info> extractInfos(
+            ChannelState handle, Comparator<Info> comparator) {
+        if (handle instanceof AbstractChannelStateHandle) {
+            return Collections.singletonList(((AbstractChannelStateHandle<Info>) handle).getInfo());
+        } else {
+            throw new IllegalStateException(
+                    "Not supported InputStateHandle : " + handle.getClass());
+        }
+    }
+
     private static <T, E> BiFunction<T, T, Boolean> eqStateApprover(
             Function<T, E> identityExtractor) {
         return (ref, alt) -> identityExtractor.apply(ref).equals(identityExtractor.apply(alt));
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 91085856ad4..6ce8d6a5932 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
@@ -34,6 +34,7 @@ import org.apache.flink.runtime.jobgraph.JobEdge;
 import org.apache.flink.runtime.jobgraph.OperatorID;
 import org.apache.flink.runtime.jobgraph.OperatorInstanceID;
 import org.apache.flink.runtime.state.AbstractChannelStateHandle;
+import org.apache.flink.runtime.state.ChannelStateHelper;
 import org.apache.flink.runtime.state.InputChannelStateHandle;
 import org.apache.flink.runtime.state.KeyGroupRange;
 import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
@@ -366,12 +367,12 @@ public class StateAssignmentOperation {
 
         checkForUnsupportedToplogyChanges(
                 assignment.oldState,
-                OperatorSubtaskState::getResultSubpartitionState,
+                ChannelStateHelper::extractUnmergedOutputHandles,
                 assignment.outputOperatorID);
 
         final OperatorState outputState = assignment.oldState.get(assignment.outputOperatorID);
         final List<List<ResultSubpartitionStateHandle>> outputOperatorState =
-                splitBySubtasks(outputState, OperatorSubtaskState::getResultSubpartitionState);
+                splitBySubtasks(outputState, ChannelStateHelper::extractUnmergedOutputHandles);
 
         final ExecutionJobVertex executionJobVertex = assignment.executionJobVertex;
         final List<IntermediateDataSet> outputs =
@@ -415,7 +416,7 @@ public class StateAssignmentOperation {
 
         checkForUnsupportedToplogyChanges(
                 stateAssignment.oldState,
-                OperatorSubtaskState::getInputChannelState,
+                ChannelStateHelper::extractUnmergedInputHandles,
                 stateAssignment.inputOperatorID);
 
         final ExecutionJobVertex executionJobVertex = stateAssignment.executionJobVertex;
@@ -425,7 +426,7 @@ public class StateAssignmentOperation {
         final OperatorState inputState =
                 stateAssignment.oldState.get(stateAssignment.inputOperatorID);
         final List<List<InputChannelStateHandle>> inputOperatorState =
-                splitBySubtasks(inputState, OperatorSubtaskState::getInputChannelState);
+                splitBySubtasks(inputState, ChannelStateHelper::extractUnmergedInputHandles);
 
         boolean hasAnyFullMapper =
                 executionJobVertex.getJobVertex().getInputs().stream()
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 c61e30407a3..75e11d147a4 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
@@ -51,6 +51,8 @@ import java.util.function.Function;
 import java.util.stream.IntStream;
 
 import static java.util.Collections.emptySet;
+import static org.apache.flink.runtime.state.ChannelStateHelper.castToInputStateCollection;
+import static org.apache.flink.runtime.state.ChannelStateHelper.castToOutputStateCollection;
 import static org.apache.flink.util.Preconditions.checkNotNull;
 import static org.apache.flink.util.Preconditions.checkState;
 
@@ -168,17 +170,15 @@ class TaskStateAssignment {
                         || !subRawKeyedState.containsKey(instanceID),
                 "If an operator has no managed key state, it should also not have a raw keyed state.");
 
-        final StateObjectCollection<InputChannelStateHandle> inputState =
-                getState(instanceID, inputChannelStates);
-        final StateObjectCollection<ResultSubpartitionStateHandle> outputState =
-                getState(instanceID, resultSubpartitionStates);
         return OperatorSubtaskState.builder()
                 .setManagedOperatorState(getState(instanceID, subManagedOperatorState))
                 .setRawOperatorState(getState(instanceID, subRawOperatorState))
                 .setManagedKeyedState(getState(instanceID, subManagedKeyedState))
                 .setRawKeyedState(getState(instanceID, subRawKeyedState))
-                .setInputChannelState(inputState)
-                .setResultSubpartitionState(outputState)
+                .setInputChannelState(
+                        castToInputStateCollection(inputChannelStates.get(instanceID)))
+                .setResultSubpartitionState(
+                        castToOutputStateCollection(resultSubpartitionStates.get(instanceID)))
                 .setInputRescalingDescriptor(
                         createRescalingDescriptor(
                                 instanceID,
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java
index 88b23c3bb7a..6fe1f1a6e99 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java
@@ -26,6 +26,7 @@ import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
 import org.apache.flink.runtime.io.network.logger.NetworkActionsLogger;
 import org.apache.flink.runtime.io.network.partition.consumer.InputGate;
 import org.apache.flink.runtime.state.AbstractChannelStateHandle;
+import org.apache.flink.runtime.state.ChannelStateHelper;
 import org.apache.flink.runtime.state.StreamStateHandle;
 
 import java.io.Closeable;
@@ -64,7 +65,8 @@ public class SequentialChannelStateReaderImpl implements SequentialChannelStateR
             read(
                     stateHandler,
                     groupByDelegate(
-                            streamSubtaskStates(), OperatorSubtaskState::getInputChannelState));
+                            streamSubtaskStates(),
+                            ChannelStateHelper::extractUnmergedInputHandles));
         }
     }
 
@@ -80,7 +82,7 @@ public class SequentialChannelStateReaderImpl implements SequentialChannelStateR
                     stateHandler,
                     groupByDelegate(
                             streamSubtaskStates(),
-                            OperatorSubtaskState::getResultSubpartitionState));
+                            ChannelStateHelper::extractUnmergedOutputHandles));
         }
     }
 
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializer.java
index 1e1979cea64..150caf1f6f4 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializer.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializer.java
@@ -17,11 +17,11 @@
 
 package org.apache.flink.runtime.checkpoint.metadata;
 
-import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo;
-import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo;
 import org.apache.flink.runtime.state.AbstractChannelStateHandle;
 import org.apache.flink.runtime.state.AbstractChannelStateHandle.StateContentMetaInfo;
 import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.InputStateHandle;
+import org.apache.flink.runtime.state.OutputStateHandle;
 import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
 import org.apache.flink.runtime.state.StreamStateHandle;
 import org.apache.flink.util.function.BiConsumerWithException;
@@ -36,50 +36,51 @@ import java.util.List;
 
 import static org.apache.flink.runtime.checkpoint.metadata.MetadataV2V3SerializerBase.deserializeStreamStateHandle;
 import static org.apache.flink.runtime.checkpoint.metadata.MetadataV2V3SerializerBase.serializeStreamStateHandle;
+import static org.apache.flink.runtime.state.ChannelStateHelper.INPUT_CHANNEL_INFO_READER;
+import static org.apache.flink.runtime.state.ChannelStateHelper.INPUT_CHANNEL_INFO_WRITER;
+import static org.apache.flink.runtime.state.ChannelStateHelper.RESULT_SUBPARTITION_INFO_READER;
+import static org.apache.flink.runtime.state.ChannelStateHelper.RESULT_SUBPARTITION_INFO_WRITER;
 
+/**
+ * ChannelStateHandlerSerializer only support {@link InputChannelStateHandle} and {@link
+ * ResultSubpartitionStateHandle}, used in {@link MetadataV3Serializer} and {@link
+ * MetadataV4Serializer}.
+ */
 class ChannelStateHandleSerializer {
 
-    public void serialize(ResultSubpartitionStateHandle handle, DataOutputStream dataOutputStream)
-            throws IOException {
+    public void serialize(OutputStateHandle handle, DataOutputStream dos) throws IOException {
+        if (!(handle instanceof ResultSubpartitionStateHandle)) {
+            throw new IllegalStateException(
+                    "OutputStateHandle must be ResultSubpartitionStateHandle.");
+        }
+
         serializeChannelStateHandle(
-                handle,
-                dataOutputStream,
-                (info, out) -> {
-                    out.writeInt(info.getPartitionIdx());
-                    out.writeInt(info.getSubPartitionIdx());
-                });
+                (ResultSubpartitionStateHandle) handle, dos, RESULT_SUBPARTITION_INFO_WRITER);
     }
 
-    ResultSubpartitionStateHandle deserializeResultSubpartitionStateHandle(
+    public OutputStateHandle deserializeOutputStateHandle(
             DataInputStream dis, MetadataV2V3SerializerBase.DeserializationContext context)
             throws IOException {
 
         return deserializeChannelStateHandle(
-                is -> new ResultSubpartitionInfo(is.readInt(), is.readInt()),
-                ResultSubpartitionStateHandle::new,
-                dis,
-                context);
+                RESULT_SUBPARTITION_INFO_READER, ResultSubpartitionStateHandle::new, dis, context);
     }
 
-    public void serialize(InputChannelStateHandle handle, DataOutputStream dos) throws IOException {
+    public void serialize(InputStateHandle handle, DataOutputStream dos) throws IOException {
+        if (!(handle instanceof InputChannelStateHandle)) {
+            throw new IllegalStateException("InputStateHandle must be InputChannelStateHandle.");
+        }
+
         serializeChannelStateHandle(
-                handle,
-                dos,
-                (info, dataOutputStream) -> {
-                    dos.writeInt(info.getGateIdx());
-                    dos.writeInt(info.getInputChannelIdx());
-                });
+                (InputChannelStateHandle) handle, dos, INPUT_CHANNEL_INFO_WRITER);
     }
 
-    InputChannelStateHandle deserializeInputChannelStateHandle(
+    public InputStateHandle deserializeInputStateHandle(
             DataInputStream dis, MetadataV2V3SerializerBase.DeserializationContext context)
             throws IOException {
 
         return deserializeChannelStateHandle(
-                is -> new InputChannelInfo(is.readInt(), is.readInt()),
-                InputChannelStateHandle::new,
-                dis,
-                context);
+                INPUT_CHANNEL_INFO_READER, InputChannelStateHandle::new, dis, context);
     }
 
     private static <I> void serializeChannelStateHandle(
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV2V3SerializerBase.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV2V3SerializerBase.java
index 9f26090fd2c..37d043047d5 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV2V3SerializerBase.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV2V3SerializerBase.java
@@ -30,7 +30,7 @@ import org.apache.flink.runtime.checkpoint.filemerging.LogicalFile;
 import org.apache.flink.runtime.state.CheckpointedStateScope;
 import org.apache.flink.runtime.state.IncrementalKeyedStateHandle.HandleAndLocalPath;
 import org.apache.flink.runtime.state.IncrementalRemoteKeyedStateHandle;
-import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.InputStateHandle;
 import org.apache.flink.runtime.state.KeyGroupRange;
 import org.apache.flink.runtime.state.KeyGroupRangeOffsets;
 import org.apache.flink.runtime.state.KeyGroupsSavepointStateHandle;
@@ -38,7 +38,7 @@ import org.apache.flink.runtime.state.KeyGroupsStateHandle;
 import org.apache.flink.runtime.state.KeyedStateHandle;
 import org.apache.flink.runtime.state.OperatorStateHandle;
 import org.apache.flink.runtime.state.OperatorStreamStateHandle;
-import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
+import org.apache.flink.runtime.state.OutputStateHandle;
 import org.apache.flink.runtime.state.StateHandleID;
 import org.apache.flink.runtime.state.StateObject;
 import org.apache.flink.runtime.state.StreamStateHandle;
@@ -696,25 +696,22 @@ public abstract class MetadataV2V3SerializerBase {
     //  channel state (unaligned checkpoints)
     // ------------------------------------------------------------------------
 
-    protected StateObjectCollection<ResultSubpartitionStateHandle>
-            deserializeResultSubpartitionStateHandle(
-                    DataInputStream dis, @Nullable DeserializationContext context)
-                    throws IOException {
+    protected StateObjectCollection<OutputStateHandle> deserializeResultSubpartitionStateHandle(
+            DataInputStream dis, @Nullable DeserializationContext context) throws IOException {
         return StateObjectCollection.empty();
     }
 
-    protected StateObjectCollection<InputChannelStateHandle> deserializeInputChannelStateHandle(
+    protected StateObjectCollection<InputStateHandle> deserializeInputChannelStateHandle(
             DataInputStream dis, @Nullable DeserializationContext context) throws IOException {
         return StateObjectCollection.empty();
     }
 
     protected void serializeResultSubpartitionStateHandle(
-            ResultSubpartitionStateHandle resultSubpartitionStateHandle, DataOutputStream dos)
+            OutputStateHandle resultSubpartitionStateHandle, DataOutputStream dos)
             throws IOException {}
 
     protected void serializeInputChannelStateHandle(
-            InputChannelStateHandle inputChannelStateHandle, DataOutputStream dos)
-            throws IOException {}
+            InputStateHandle inputChannelStateHandle, DataOutputStream dos) throws IOException {}
 
     // ------------------------------------------------------------------------
     //  low-level state handles
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV3Serializer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV3Serializer.java
index fc6df97ce8d..c2fe5134e08 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV3Serializer.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/MetadataV3Serializer.java
@@ -26,10 +26,10 @@ import org.apache.flink.runtime.checkpoint.OperatorState;
 import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
 import org.apache.flink.runtime.checkpoint.StateObjectCollection;
 import org.apache.flink.runtime.jobgraph.OperatorID;
-import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.InputStateHandle;
 import org.apache.flink.runtime.state.KeyedStateHandle;
 import org.apache.flink.runtime.state.OperatorStateHandle;
-import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
+import org.apache.flink.runtime.state.OutputStateHandle;
 import org.apache.flink.runtime.state.StateObject;
 import org.apache.flink.runtime.state.StreamStateHandle;
 import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
@@ -203,35 +203,31 @@ public class MetadataV3Serializer extends MetadataV2V3SerializerBase implements
     @VisibleForTesting
     @Override
     public void serializeResultSubpartitionStateHandle(
-            ResultSubpartitionStateHandle handle, DataOutputStream dos) throws IOException {
+            OutputStateHandle handle, DataOutputStream dos) throws IOException {
         channelStateHandleSerializer.serialize(handle, dos);
     }
 
     @VisibleForTesting
     @Override
-    public StateObjectCollection<ResultSubpartitionStateHandle>
-            deserializeResultSubpartitionStateHandle(
-                    DataInputStream dis, @Nullable DeserializationContext context)
-                    throws IOException {
+    public StateObjectCollection<OutputStateHandle> deserializeResultSubpartitionStateHandle(
+            DataInputStream dis, @Nullable DeserializationContext context) throws IOException {
         return deserializeCollection(
-                dis,
-                context,
-                channelStateHandleSerializer::deserializeResultSubpartitionStateHandle);
+                dis, context, channelStateHandleSerializer::deserializeOutputStateHandle);
     }
 
     @VisibleForTesting
     @Override
-    public void serializeInputChannelStateHandle(
-            InputChannelStateHandle handle, DataOutputStream dos) throws IOException {
+    public void serializeInputChannelStateHandle(InputStateHandle handle, DataOutputStream dos)
+            throws IOException {
         channelStateHandleSerializer.serialize(handle, dos);
     }
 
     @VisibleForTesting
     @Override
-    public StateObjectCollection<InputChannelStateHandle> deserializeInputChannelStateHandle(
+    public StateObjectCollection<InputStateHandle> deserializeInputChannelStateHandle(
             DataInputStream dis, @Nullable DeserializationContext context) throws IOException {
         return deserializeCollection(
-                dis, context, channelStateHandleSerializer::deserializeInputChannelStateHandle);
+                dis, context, channelStateHandleSerializer::deserializeInputStateHandle);
     }
 
     private <T extends StateObject> void serializeCollection(
@@ -296,14 +292,14 @@ public class MetadataV3Serializer extends MetadataV2V3SerializerBase implements
     }
 
     @VisibleForTesting
-    public static StateObjectCollection<InputChannelStateHandle> deserializeInputChannelStateHandle(
+    public static StateObjectCollection<InputStateHandle> deserializeInputChannelStateHandle(
             DataInputStream dis) throws IOException {
         return INSTANCE.deserializeInputChannelStateHandle(dis, null);
     }
 
     @VisibleForTesting
-    public StateObjectCollection<ResultSubpartitionStateHandle>
-            deserializeResultSubpartitionStateHandle(DataInputStream dis) throws IOException {
+    public StateObjectCollection<OutputStateHandle> deserializeResultSubpartitionStateHandle(
+            DataInputStream dis) throws IOException {
         return INSTANCE.deserializeResultSubpartitionStateHandle(dis, null);
     }
 
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/ChannelStateHelper.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/ChannelStateHelper.java
new file mode 100644
index 00000000000..ff72654ea31
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/ChannelStateHelper.java
@@ -0,0 +1,140 @@
+/*
+ * 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.state;
+
+import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
+import org.apache.flink.runtime.checkpoint.StateObjectCollection;
+import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo;
+import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo;
+import org.apache.flink.util.function.BiConsumerWithException;
+import org.apache.flink.util.function.FunctionWithException;
+
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.util.Collection;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+/**
+ * Utility class for channel info serialization and conversions between single channel state handle
+ * and merged state handle.
+ */
+public final class ChannelStateHelper {
+
+    private ChannelStateHelper() {}
+
+    public static final BiConsumerWithException<
+                    ResultSubpartitionInfo, DataOutputStream, IOException>
+            RESULT_SUBPARTITION_INFO_WRITER =
+                    (info, out) -> {
+                        out.writeInt(info.getPartitionIdx());
+                        out.writeInt(info.getSubPartitionIdx());
+                    };
+
+    public static final BiConsumerWithException<InputChannelInfo, DataOutputStream, IOException>
+            INPUT_CHANNEL_INFO_WRITER =
+                    (info, dos) -> {
+                        dos.writeInt(info.getGateIdx());
+                        dos.writeInt(info.getInputChannelIdx());
+                    };
+
+    public static final FunctionWithException<DataInputStream, ResultSubpartitionInfo, IOException>
+            RESULT_SUBPARTITION_INFO_READER =
+                    is -> new ResultSubpartitionInfo(is.readInt(), is.readInt());
+
+    public static final FunctionWithException<DataInputStream, InputChannelInfo, IOException>
+            INPUT_CHANNEL_INFO_READER = is -> new InputChannelInfo(is.readInt(), is.readInt());
+
+    public static Stream<StreamStateHandle> collectUniqueDisposableInChannelState(
+            Stream<StateObjectCollection<? extends ChannelState>> stateCollections) {
+        return stateCollections
+                .flatMap(Collection::stream)
+                .map(
+                        handle -> {
+                            if (handle instanceof AbstractChannelStateHandle) {
+                                return ((AbstractChannelStateHandle<?>) handle).getDelegate();
+
+                            } else {
+                                throw new IllegalStateException(
+                                        "Not Supported state handle : " + handle.getClass());
+                            }
+                        })
+                .distinct();
+    }
+
+    public static StateObjectCollection<InputStateHandle> castToInputStateCollection(
+            Collection<InputChannelStateHandle> handles) {
+        if (handles == null || handles.isEmpty()) {
+            return StateObjectCollection.empty();
+        }
+
+        return new StateObjectCollection<>(
+                handles.stream().map(e -> ((InputStateHandle) e)).collect(Collectors.toList()));
+    }
+
+    public static StateObjectCollection<OutputStateHandle> castToOutputStateCollection(
+            Collection<ResultSubpartitionStateHandle> handles) {
+        if (handles == null || handles.isEmpty()) {
+            return StateObjectCollection.empty();
+        }
+
+        return new StateObjectCollection<>(
+                handles.stream().map(e -> ((OutputStateHandle) e)).collect(Collectors.toList()));
+    }
+
+    public static StateObjectCollection<InputChannelStateHandle> extractUnmergedInputHandles(
+            OperatorSubtaskState subtaskState) {
+        List<InputChannelStateHandle> inputHandles =
+                subtaskState.getInputChannelState().stream()
+                        .flatMap(
+                                h -> {
+                                    if (h instanceof InputChannelStateHandle) {
+                                        return Stream.of((InputChannelStateHandle) h);
+                                    } else {
+                                        throw new IllegalStateException(
+                                                "Invalid input channel state : " + h.getClass());
+                                    }
+                                })
+                        .collect(Collectors.toList());
+        return inputHandles.isEmpty()
+                ? StateObjectCollection.empty()
+                : new StateObjectCollection<>(inputHandles);
+    }
+
+    public static StateObjectCollection<ResultSubpartitionStateHandle> extractUnmergedOutputHandles(
+            OperatorSubtaskState subtaskState) {
+        List<ResultSubpartitionStateHandle> outputHandles =
+                subtaskState.getResultSubpartitionState().stream()
+                        .flatMap(
+                                h -> {
+                                    if (h instanceof ResultSubpartitionStateHandle) {
+                                        return Stream.of((ResultSubpartitionStateHandle) h);
+                                    } else {
+                                        throw new IllegalStateException(
+                                                "Invalid output channel state : " + h.getClass());
+                                    }
+                                })
+                        .collect(Collectors.toList());
+        return outputHandles.isEmpty()
+                ? StateObjectCollection.empty()
+                : new StateObjectCollection<>(outputHandles);
+    }
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFinalizer.java b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFinalizer.java
index 27051f1d57c..c308a6cb8ae 100644
--- a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFinalizer.java
+++ b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFinalizer.java
@@ -20,10 +20,10 @@ package org.apache.flink.streaming.api.operators;
 
 import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
 import org.apache.flink.runtime.checkpoint.StateObjectCollection;
-import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.InputStateHandle;
 import org.apache.flink.runtime.state.KeyedStateHandle;
 import org.apache.flink.runtime.state.OperatorStateHandle;
-import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
+import org.apache.flink.runtime.state.OutputStateHandle;
 import org.apache.flink.runtime.state.SnapshotResult;
 import org.apache.flink.util.concurrent.FutureUtils;
 
@@ -62,10 +62,10 @@ public class OperatorSnapshotFinalizer {
         SnapshotResult<OperatorStateHandle> operatorRaw =
                 FutureUtils.runIfNotDoneAndGet(snapshotFutures.getOperatorStateRawFuture());
 
-        SnapshotResult<StateObjectCollection<InputChannelStateHandle>> inputChannel =
+        SnapshotResult<StateObjectCollection<InputStateHandle>> inputChannel =
                 snapshotFutures.getInputChannelStateFuture().get();
 
-        SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>> resultSubpartition =
+        SnapshotResult<StateObjectCollection<OutputStateHandle>> resultSubpartition =
                 snapshotFutures.getResultSubpartitionStateFuture().get();
 
         jobManagerOwnedState =
diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFutures.java b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFutures.java
index a8dad0f5de2..ce36d1fe5ba 100644
--- a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFutures.java
+++ b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFutures.java
@@ -21,10 +21,10 @@ package org.apache.flink.streaming.api.operators;
 import org.apache.flink.api.java.tuple.Tuple2;
 import org.apache.flink.runtime.checkpoint.StateObjectCollection;
 import org.apache.flink.runtime.state.DoneFuture;
-import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.InputStateHandle;
 import org.apache.flink.runtime.state.KeyedStateHandle;
 import org.apache.flink.runtime.state.OperatorStateHandle;
-import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
+import org.apache.flink.runtime.state.OutputStateHandle;
 import org.apache.flink.runtime.state.SnapshotResult;
 import org.apache.flink.runtime.state.StateObject;
 
@@ -51,11 +51,10 @@ public class OperatorSnapshotFutures {
     @Nonnull private RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateRawFuture;
 
     @Nonnull
-    private Future<SnapshotResult<StateObjectCollection<InputChannelStateHandle>>>
-            inputChannelStateFuture;
+    private Future<SnapshotResult<StateObjectCollection<InputStateHandle>>> inputChannelStateFuture;
 
     @Nonnull
-    private Future<SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>>>
+    private Future<SnapshotResult<StateObjectCollection<OutputStateHandle>>>
             resultSubpartitionStateFuture;
 
     public OperatorSnapshotFutures() {
@@ -74,10 +73,10 @@ public class OperatorSnapshotFutures {
             @Nonnull RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateManagedFuture,
             @Nonnull RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateRawFuture,
             @Nonnull
-                    Future<SnapshotResult<StateObjectCollection<InputChannelStateHandle>>>
+                    Future<SnapshotResult<StateObjectCollection<InputStateHandle>>>
                             inputChannelStateFuture,
             @Nonnull
-                    Future<SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>>>
+                    Future<SnapshotResult<StateObjectCollection<OutputStateHandle>>>
                             resultSubpartitionStateFuture) {
         this.keyedStateManagedFuture = keyedStateManagedFuture;
         this.keyedStateRawFuture = keyedStateRawFuture;
@@ -130,27 +129,27 @@ public class OperatorSnapshotFutures {
     }
 
     @Nonnull
-    public Future<SnapshotResult<StateObjectCollection<InputChannelStateHandle>>>
+    public Future<SnapshotResult<StateObjectCollection<InputStateHandle>>>
             getInputChannelStateFuture() {
         return inputChannelStateFuture;
     }
 
     public void setInputChannelStateFuture(
             @Nonnull
-                    Future<SnapshotResult<StateObjectCollection<InputChannelStateHandle>>>
+                    Future<SnapshotResult<StateObjectCollection<InputStateHandle>>>
                             inputChannelStateFuture) {
         this.inputChannelStateFuture = inputChannelStateFuture;
     }
 
     @Nonnull
-    public Future<SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>>>
+    public Future<SnapshotResult<StateObjectCollection<OutputStateHandle>>>
             getResultSubpartitionStateFuture() {
         return resultSubpartitionStateFuture;
     }
 
     public void setResultSubpartitionStateFuture(
             @Nonnull
-                    Future<SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>>>
+                    Future<SnapshotResult<StateObjectCollection<OutputStateHandle>>>
                             resultSubpartitionStateFuture) {
         this.resultSubpartitionStateFuture = resultSubpartitionStateFuture;
     }
diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/OperatorChain.java b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/OperatorChain.java
index 7e17d63b183..febbcb1bdb0 100644
--- a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/OperatorChain.java
+++ b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/OperatorChain.java
@@ -469,14 +469,14 @@ public abstract class OperatorChain<OUT, OP extends StreamOperator<OUT>>
             snapshotInProgress.setInputChannelStateFuture(
                     channelStateWriteResult
                             .getInputChannelStateHandles()
-                            .thenApply(StateObjectCollection::new)
+                            .thenApply(e -> new StateObjectCollection(e))
                             .thenApply(SnapshotResult::of));
         }
         if (op == getTailOperator()) {
             snapshotInProgress.setResultSubpartitionStateFuture(
                     channelStateWriteResult
                             .getResultSubpartitionStateHandles()
-                            .thenApply(StateObjectCollection::new)
+                            .thenApply(e -> new StateObjectCollection(e))
                             .thenApply(SnapshotResult::of));
         }
     }
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorFailureTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorFailureTest.java
index 52cc0d1bcf2..f10d92d1b86 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorFailureTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorFailureTest.java
@@ -53,9 +53,11 @@ import java.util.List;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Stream;
 
 import static java.util.Collections.emptyList;
 import static org.apache.flink.runtime.checkpoint.CheckpointCoordinatorTest.assertStatsMetrics;
+import static org.apache.flink.runtime.state.ChannelStateHelper.collectUniqueDisposableInChannelState;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.fail;
 import static org.mockito.Mockito.mock;
@@ -174,9 +176,18 @@ class CheckpointCoordinatorFailureTest {
         verify(operatorSubtaskState.getRawOperatorState().iterator().next()).discardState();
         verify(operatorSubtaskState.getManagedKeyedState().iterator().next()).discardState();
         verify(operatorSubtaskState.getRawKeyedState().iterator().next()).discardState();
-        verify(operatorSubtaskState.getInputChannelState().iterator().next().getDelegate())
+        verify(
+                        collectUniqueDisposableInChannelState(
+                                        Stream.of(operatorSubtaskState.getInputChannelState()))
+                                .iterator()
+                                .next())
                 .discardState();
-        verify(operatorSubtaskState.getResultSubpartitionState().iterator().next().getDelegate())
+        verify(
+                        collectUniqueDisposableInChannelState(
+                                        Stream.of(
+                                                operatorSubtaskState.getResultSubpartitionState()))
+                                .iterator()
+                                .next())
                 .discardState();
     }
 
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImplTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImplTest.java
index 8aeec586b75..5f05af0d92e 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImplTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImplTest.java
@@ -20,7 +20,6 @@ package org.apache.flink.runtime.checkpoint.channel;
 import org.apache.flink.api.java.tuple.Tuple2;
 import org.apache.flink.core.memory.MemorySegmentFactory;
 import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
-import org.apache.flink.runtime.checkpoint.StateObjectCollection;
 import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
 import org.apache.flink.runtime.io.network.buffer.Buffer;
 import org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler;
@@ -70,6 +69,8 @@ import static java.util.function.Function.identity;
 import static java.util.stream.Collectors.toList;
 import static java.util.stream.Collectors.toMap;
 import static java.util.stream.IntStream.range;
+import static org.apache.flink.runtime.state.ChannelStateHelper.castToInputStateCollection;
+import static org.apache.flink.runtime.state.ChannelStateHelper.castToOutputStateCollection;
 import static org.assertj.core.api.Assertions.assertThat;
 
 /** {@link SequentialChannelStateReaderImpl} Test. */
@@ -268,8 +269,8 @@ public class SequentialChannelStateReaderImplTest {
                 Collections.singletonMap(
                         new OperatorID(),
                         OperatorSubtaskState.builder()
-                                .setInputChannelState(new StateObjectCollection<>(handles.f0))
-                                .setResultSubpartitionState(new StateObjectCollection<>(handles.f1))
+                                .setInputChannelState(castToInputStateCollection(handles.f0))
+                                .setResultSubpartitionState(castToOutputStateCollection(handles.f1))
                                 .build()));
     }
 
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/ChannelPersistenceITCase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/ChannelPersistenceITCase.java
index 30b9aed26af..53c05b19613 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/ChannelPersistenceITCase.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/ChannelPersistenceITCase.java
@@ -21,7 +21,6 @@ import org.apache.flink.api.common.JobID;
 import org.apache.flink.core.memory.MemorySegmentFactory;
 import org.apache.flink.runtime.checkpoint.CheckpointOptions;
 import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
-import org.apache.flink.runtime.checkpoint.StateObjectCollection;
 import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
 import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriteRequestExecutorFactory;
 import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter.ChannelStateWriteResult;
@@ -71,6 +70,8 @@ import static java.util.Collections.singletonMap;
 import static org.apache.flink.runtime.checkpoint.CheckpointType.CHECKPOINT;
 import static org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter.SEQUENCE_NUMBER_UNKNOWN;
 import static org.apache.flink.runtime.io.network.buffer.Buffer.DataType.RECOVERY_COMPLETION;
+import static org.apache.flink.runtime.state.ChannelStateHelper.castToInputStateCollection;
+import static org.apache.flink.runtime.state.ChannelStateHelper.castToOutputStateCollection;
 import static org.apache.flink.util.CloseableIterator.ofElements;
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -305,10 +306,10 @@ class ChannelPersistenceITCase {
                         new OperatorID(),
                         OperatorSubtaskState.builder()
                                 .setInputChannelState(
-                                        new StateObjectCollection<>(
+                                        castToInputStateCollection(
                                                 t.getInputChannelStateHandles().get()))
                                 .setResultSubpartitionState(
-                                        new StateObjectCollection<>(
+                                        castToOutputStateCollection(
                                                 t.getResultSubpartitionStateHandles().get()))
                                 .build()));
     }
diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFinalizerTest.java b/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFinalizerTest.java
index f51fdcff3af..a64905098af 100644
--- a/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFinalizerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFinalizerTest.java
@@ -23,9 +23,11 @@ import org.apache.flink.runtime.checkpoint.StateHandleDummyUtil;
 import org.apache.flink.runtime.checkpoint.StateObjectCollection;
 import org.apache.flink.runtime.state.DoneFuture;
 import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.InputStateHandle;
 import org.apache.flink.runtime.state.KeyGroupRange;
 import org.apache.flink.runtime.state.KeyedStateHandle;
 import org.apache.flink.runtime.state.OperatorStateHandle;
+import org.apache.flink.runtime.state.OutputStateHandle;
 import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
 import org.apache.flink.runtime.state.SnapshotResult;
 import org.apache.flink.runtime.state.StateObject;
@@ -69,11 +71,11 @@ class OperatorSnapshotFinalizerTest {
                 withLocalState(deepDummyCopy(operatorTemplate), deepDummyCopy(operatorTemplate));
         SnapshotResult<OperatorStateHandle> rawOper =
                 withLocalState(deepDummyCopy(operatorTemplate), deepDummyCopy(operatorTemplate));
-        SnapshotResult<StateObjectCollection<InputChannelStateHandle>> inputChannel =
+        SnapshotResult<StateObjectCollection<InputStateHandle>> inputChannel =
                 withLocalState(
                         singleton(deepDummyCopy(inputChannelTemplate)),
                         singleton(deepDummyCopy(inputChannelTemplate)));
-        SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>> resultSubpartition =
+        SnapshotResult<StateObjectCollection<OutputStateHandle>> resultSubpartition =
                 withLocalState(
                         singleton(deepDummyCopy(resultSubpartitionTemplate)),
                         singleton(deepDummyCopy(resultSubpartitionTemplate)));
diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFuturesTest.java b/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFuturesTest.java
index e984616468b..d9c6c23283b 100644
--- a/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFuturesTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFuturesTest.java
@@ -22,11 +22,13 @@ import org.apache.flink.api.java.tuple.Tuple2;
 import org.apache.flink.runtime.checkpoint.StateObjectCollection;
 import org.apache.flink.runtime.state.DoneFuture;
 import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.InputStateHandle;
 import org.apache.flink.runtime.state.KeyGroupRangeOffsets;
 import org.apache.flink.runtime.state.KeyGroupsStateHandle;
 import org.apache.flink.runtime.state.KeyedStateHandle;
 import org.apache.flink.runtime.state.OperatorStateHandle;
 import org.apache.flink.runtime.state.OperatorStreamStateHandle;
+import org.apache.flink.runtime.state.OutputStateHandle;
 import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
 import org.apache.flink.runtime.state.SnapshotResult;
 import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
@@ -102,18 +104,17 @@ class OperatorSnapshotFuturesTest {
                 spy(DoneFuture.of(operatorStateRawResult));
 
         InputChannelStateHandle inputChannelRawStateHandle = mock(InputChannelStateHandle.class);
-        SnapshotResult<StateObjectCollection<InputChannelStateHandle>> inputChannelStateRawResult =
+        SnapshotResult<StateObjectCollection<InputStateHandle>> inputChannelStateRawResult =
                 SnapshotResult.of(StateObjectCollection.singleton(inputChannelRawStateHandle));
-        Future<SnapshotResult<StateObjectCollection<InputChannelStateHandle>>>
-                inputChannelStateRawFuture = spy(DoneFuture.of(inputChannelStateRawResult));
+        Future<SnapshotResult<StateObjectCollection<InputStateHandle>>> inputChannelStateRawFuture =
+                spy(DoneFuture.of(inputChannelStateRawResult));
 
         ResultSubpartitionStateHandle resultSubpartitionRawStateHandle =
                 mock(ResultSubpartitionStateHandle.class);
-        SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>>
-                resultSubpartitionStateRawResult =
-                        SnapshotResult.of(
-                                StateObjectCollection.singleton(resultSubpartitionRawStateHandle));
-        Future<SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>>>
+        SnapshotResult<StateObjectCollection<OutputStateHandle>> resultSubpartitionStateRawResult =
+                SnapshotResult.of(
+                        StateObjectCollection.singleton(resultSubpartitionRawStateHandle));
+        Future<SnapshotResult<StateObjectCollection<OutputStateHandle>>>
                 resultSubpartitionStateRawFuture =
                         spy(DoneFuture.of(resultSubpartitionStateRawResult));
 
diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/StreamOperatorStateHandlerTest.java b/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/StreamOperatorStateHandlerTest.java
index 9cb1d4522f2..2b827aff041 100644
--- a/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/StreamOperatorStateHandlerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/StreamOperatorStateHandlerTest.java
@@ -32,11 +32,11 @@ import org.apache.flink.runtime.metrics.util.InterceptingOperatorMetricGroup;
 import org.apache.flink.runtime.operators.testutils.ExpectedTestException;
 import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
 import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
-import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.InputStateHandle;
 import org.apache.flink.runtime.state.KeyGroupRange;
 import org.apache.flink.runtime.state.KeyedStateHandle;
 import org.apache.flink.runtime.state.OperatorStateHandle;
-import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
+import org.apache.flink.runtime.state.OutputStateHandle;
 import org.apache.flink.runtime.state.SnapshotResult;
 import org.apache.flink.runtime.state.StateInitializationContext;
 import org.apache.flink.runtime.state.StateSnapshotContext;
@@ -77,9 +77,9 @@ class StreamOperatorStateHandlerTest {
                     new CancelableFuture<>();
             RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateRawFuture =
                     new CancelableFuture<>();
-            RunnableFuture<SnapshotResult<StateObjectCollection<InputChannelStateHandle>>>
+            RunnableFuture<SnapshotResult<StateObjectCollection<InputStateHandle>>>
                     inputChannelStateFuture = new CancelableFuture<>();
-            RunnableFuture<SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>>>
+            RunnableFuture<SnapshotResult<StateObjectCollection<OutputStateHandle>>>
                     resultSubpartitionStateFuture = new CancelableFuture<>();
 
             OperatorSnapshotFutures operatorSnapshotResult =
diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/util/OperatorSnapshotUtil.java b/flink-runtime/src/test/java/org/apache/flink/streaming/util/OperatorSnapshotUtil.java
index 150ef6b8edb..7f9312f0c38 100644
--- a/flink-runtime/src/test/java/org/apache/flink/streaming/util/OperatorSnapshotUtil.java
+++ b/flink-runtime/src/test/java/org/apache/flink/streaming/util/OperatorSnapshotUtil.java
@@ -21,10 +21,10 @@ package org.apache.flink.streaming.util;
 import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
 import org.apache.flink.runtime.checkpoint.StateObjectCollection;
 import org.apache.flink.runtime.checkpoint.metadata.MetadataV3Serializer;
-import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.InputStateHandle;
 import org.apache.flink.runtime.state.KeyedStateHandle;
 import org.apache.flink.runtime.state.OperatorStateHandle;
-import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
+import org.apache.flink.runtime.state.OutputStateHandle;
 
 import java.io.DataInputStream;
 import java.io.DataOutputStream;
@@ -101,19 +101,17 @@ public class OperatorSnapshotUtil {
                 dos.writeInt(-1);
             }
 
-            Collection<InputChannelStateHandle> inputChannelStateHandles =
-                    state.getInputChannelState();
+            Collection<InputStateHandle> inputChannelStateHandles = state.getInputChannelState();
             dos.writeInt(inputChannelStateHandles.size());
-            for (InputChannelStateHandle inputChannelStateHandle : inputChannelStateHandles) {
+            for (InputStateHandle inputChannelStateHandle : inputChannelStateHandles) {
                 MetadataV3Serializer.INSTANCE.serializeInputChannelStateHandle(
                         inputChannelStateHandle, dos);
             }
 
-            Collection<ResultSubpartitionStateHandle> resultSubpartitionStateHandles =
+            Collection<OutputStateHandle> resultSubpartitionStateHandles =
                     state.getResultSubpartitionState();
             dos.writeInt(inputChannelStateHandles.size());
-            for (ResultSubpartitionStateHandle resultSubpartitionStateHandle :
-                    resultSubpartitionStateHandles) {
+            for (OutputStateHandle resultSubpartitionStateHandle : resultSubpartitionStateHandles) {
                 MetadataV3Serializer.INSTANCE.serializeResultSubpartitionStateHandle(
                         resultSubpartitionStateHandle, dos);
             }
@@ -177,17 +175,16 @@ public class OperatorSnapshotUtil {
                 }
             }
 
-            final StateObjectCollection<InputChannelStateHandle> inputChannelStateHandles =
+            final StateObjectCollection<InputStateHandle> inputChannelStateHandles =
                     v == MetadataV3Serializer.VERSION
                             ? MetadataV3Serializer.deserializeInputChannelStateHandle(dis)
                             : StateObjectCollection.empty();
 
-            final StateObjectCollection<ResultSubpartitionStateHandle>
-                    resultSubpartitionStateHandles =
-                            v == MetadataV3Serializer.VERSION
-                                    ? MetadataV3Serializer.INSTANCE
-                                            .deserializeResultSubpartitionStateHandle(dis)
-                                    : StateObjectCollection.empty();
+            final StateObjectCollection<OutputStateHandle> resultSubpartitionStateHandles =
+                    v == MetadataV3Serializer.VERSION
+                            ? MetadataV3Serializer.INSTANCE
+                                    .deserializeResultSubpartitionStateHandle(dis)
+                            : StateObjectCollection.empty();
 
             return OperatorSubtaskState.builder()
                     .setManagedOperatorState(new StateObjectCollection<>(managedOperatorState))
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/LocalStateForwardingTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/LocalStateForwardingTest.java
index eceda40a3f5..138e129b85c 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/LocalStateForwardingTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/LocalStateForwardingTest.java
@@ -33,12 +33,12 @@ import org.apache.flink.runtime.jobgraph.JobVertexID;
 import org.apache.flink.runtime.jobgraph.OperatorID;
 import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
 import org.apache.flink.runtime.state.DoneFuture;
-import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.InputStateHandle;
 import org.apache.flink.runtime.state.KeyedStateHandle;
 import org.apache.flink.runtime.state.LocalRecoveryConfig;
 import org.apache.flink.runtime.state.LocalSnapshotDirectoryProviderImpl;
 import org.apache.flink.runtime.state.OperatorStateHandle;
-import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
+import org.apache.flink.runtime.state.OutputStateHandle;
 import org.apache.flink.runtime.state.SnapshotResult;
 import org.apache.flink.runtime.state.StateObject;
 import org.apache.flink.runtime.state.TaskExecutorStateChangelogStoragesManager;
@@ -111,10 +111,9 @@ class LocalStateForwardingTest {
         osFuture.setKeyedStateRawFuture(createSnapshotResult(KeyedStateHandle.class));
         osFuture.setOperatorStateManagedFuture(createSnapshotResult(OperatorStateHandle.class));
         osFuture.setOperatorStateRawFuture(createSnapshotResult(OperatorStateHandle.class));
-        osFuture.setInputChannelStateFuture(
-                createSnapshotCollectionResult(InputChannelStateHandle.class));
+        osFuture.setInputChannelStateFuture(createSnapshotCollectionResult(InputStateHandle.class));
         osFuture.setResultSubpartitionStateFuture(
-                createSnapshotCollectionResult(ResultSubpartitionStateHandle.class));
+                createSnapshotCollectionResult(OutputStateHandle.class));
 
         OperatorID operatorID = new OperatorID();
         snapshots.put(operatorID, osFuture);


(flink) 04/07: [FLINK-33734][checkpointing] Introduce merged channel state handle

Posted by za...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

zakelly pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit 6f6bed67522383290cc3873dcdef46729a6815f6
Author: wangfeifan <zo...@163.com>
AuthorDate: Thu Feb 1 17:32:50 2024 +0800

    [FLINK-33734][checkpointing] Introduce merged channel state handle
---
 .../PrioritizedOperatorSubtaskState.java           |   5 +
 .../metadata/ChannelStateHandleSerializerV2.java   | 169 +++++++++++++++
 .../state/AbstractMergedChannelStateHandle.java    | 240 +++++++++++++++++++++
 .../flink/runtime/state/ChannelStateHelper.java    |  46 ++++
 .../state/MergedInputChannelStateHandle.java       |  75 +++++++
 .../state/MergedResultSubpartitionStateHandle.java |  76 +++++++
 .../ChannelStateHandleSerializerV1Test.java        |  88 +-------
 .../ChannelStateHandleSerializerV2Test.java        |  57 +++++
 .../checkpoint/metadata/ChannelStateTestUtils.java | 175 +++++++++++++++
 .../state/MergedChannelStateHandleTest.java        |  49 +++++
 10 files changed, 902 insertions(+), 78 deletions(-)

diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PrioritizedOperatorSubtaskState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PrioritizedOperatorSubtaskState.java
index 3e1aa039f0d..37d2acfe65a 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PrioritizedOperatorSubtaskState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PrioritizedOperatorSubtaskState.java
@@ -22,6 +22,7 @@ import org.apache.flink.annotation.Internal;
 import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo;
 import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo;
 import org.apache.flink.runtime.state.AbstractChannelStateHandle;
+import org.apache.flink.runtime.state.AbstractMergedChannelStateHandle;
 import org.apache.flink.runtime.state.ChannelState;
 import org.apache.flink.runtime.state.InputStateHandle;
 import org.apache.flink.runtime.state.KeyedStateHandle;
@@ -521,6 +522,10 @@ public class PrioritizedOperatorSubtaskState {
             ChannelState handle, Comparator<Info> comparator) {
         if (handle instanceof AbstractChannelStateHandle) {
             return Collections.singletonList(((AbstractChannelStateHandle<Info>) handle).getInfo());
+        } else if (handle instanceof AbstractMergedChannelStateHandle) {
+            return ((AbstractMergedChannelStateHandle<Info, AbstractChannelStateHandle<Info>>)
+                            handle)
+                    .getInfos().stream().sorted(comparator).collect(Collectors.toList());
         } else {
             throw new IllegalStateException(
                     "Not supported InputStateHandle : " + handle.getClass());
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV2.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV2.java
new file mode 100644
index 00000000000..ba8f8ea6c6c
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV2.java
@@ -0,0 +1,169 @@
+/*
+ * 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.metadata;
+
+import org.apache.flink.api.java.tuple.Tuple4;
+import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo;
+import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo;
+import org.apache.flink.runtime.state.AbstractChannelStateHandle;
+import org.apache.flink.runtime.state.AbstractMergedChannelStateHandle;
+import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.InputStateHandle;
+import org.apache.flink.runtime.state.MergedInputChannelStateHandle;
+import org.apache.flink.runtime.state.MergedResultSubpartitionStateHandle;
+import org.apache.flink.runtime.state.OutputStateHandle;
+import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+
+import static org.apache.flink.runtime.checkpoint.metadata.ChannelStateHandleSerializerV1.deserializeChannelStateHandle;
+import static org.apache.flink.runtime.checkpoint.metadata.ChannelStateHandleSerializerV1.serializeChannelStateHandle;
+import static org.apache.flink.runtime.checkpoint.metadata.MetadataV2V3SerializerBase.deserializeStreamStateHandle;
+import static org.apache.flink.runtime.checkpoint.metadata.MetadataV2V3SerializerBase.serializeStreamStateHandle;
+import static org.apache.flink.runtime.state.ChannelStateHelper.INPUT_CHANNEL_INFO_READER;
+import static org.apache.flink.runtime.state.ChannelStateHelper.INPUT_CHANNEL_INFO_WRITER;
+import static org.apache.flink.runtime.state.ChannelStateHelper.RESULT_SUBPARTITION_INFO_READER;
+import static org.apache.flink.runtime.state.ChannelStateHelper.RESULT_SUBPARTITION_INFO_WRITER;
+
+/**
+ * ChannelStateHandlerSerializer support extending handle types, first used in {@link
+ * MetadataV5Serializer} .
+ */
+class ChannelStateHandleSerializerV2 implements ChannelStateHandleSerializer {
+
+    private static final byte INPUT_CHANNEL_STATE_HANDLE = 1;
+    private static final byte RESULT_SUBPARTITION_STATE_HANDLE = 2;
+    private static final byte MERGED_INPUT_CHANNEL_STATE_HANDLE = 3;
+    private static final byte MERGED_RESULT_SUBPARTITION_STATE_HANDLE = 4;
+
+    @Override
+    public void serialize(OutputStateHandle handle, DataOutputStream dos) throws IOException {
+        if (handle instanceof ResultSubpartitionStateHandle) {
+            dos.writeByte(RESULT_SUBPARTITION_STATE_HANDLE);
+
+            serializeChannelStateHandle(
+                    (ResultSubpartitionStateHandle) handle, dos, RESULT_SUBPARTITION_INFO_WRITER);
+        } else if (handle instanceof MergedResultSubpartitionStateHandle) {
+            dos.writeByte(MERGED_RESULT_SUBPARTITION_STATE_HANDLE);
+
+            serializeMergedChannelStateHandle(
+                    (AbstractMergedChannelStateHandle<
+                                    ResultSubpartitionInfo,
+                                    AbstractChannelStateHandle<ResultSubpartitionInfo>>)
+                            handle,
+                    dos);
+        } else {
+            throw new IllegalStateException("Unsupported input state handle.");
+        }
+    }
+
+    @Override
+    public OutputStateHandle deserializeOutputStateHandle(
+            DataInputStream dis, MetadataV2V3SerializerBase.DeserializationContext context)
+            throws IOException {
+
+        final int type = dis.read();
+
+        if (RESULT_SUBPARTITION_STATE_HANDLE == type) {
+            return deserializeChannelStateHandle(
+                    RESULT_SUBPARTITION_INFO_READER,
+                    ResultSubpartitionStateHandle::new,
+                    dis,
+                    context);
+
+        } else if (MERGED_RESULT_SUBPARTITION_STATE_HANDLE == type) {
+            Tuple4<Integer, StreamStateHandle, Long, byte[]> t4 =
+                    deserializeMergedChannelStateHandle(dis, context);
+            return new MergedResultSubpartitionStateHandle(t4.f0, t4.f1, t4.f2, t4.f3);
+        } else {
+            throw new IllegalStateException("Unknown output state handle type.");
+        }
+    }
+
+    @Override
+    public void serialize(InputStateHandle handle, DataOutputStream dos) throws IOException {
+        if (handle instanceof InputChannelStateHandle) {
+            dos.writeByte(INPUT_CHANNEL_STATE_HANDLE);
+
+            serializeChannelStateHandle(
+                    (InputChannelStateHandle) handle, dos, INPUT_CHANNEL_INFO_WRITER);
+        } else if (handle instanceof MergedInputChannelStateHandle) {
+            dos.writeByte(MERGED_INPUT_CHANNEL_STATE_HANDLE);
+
+            serializeMergedChannelStateHandle(
+                    (AbstractMergedChannelStateHandle<
+                                    InputChannelInfo, AbstractChannelStateHandle<InputChannelInfo>>)
+                            handle,
+                    dos);
+        } else {
+            throw new IllegalStateException("Unsupported output state handle.");
+        }
+    }
+
+    @Override
+    public InputStateHandle deserializeInputStateHandle(
+            DataInputStream dis, MetadataV2V3SerializerBase.DeserializationContext context)
+            throws IOException {
+
+        final int type = dis.read();
+
+        if (INPUT_CHANNEL_STATE_HANDLE == type) {
+            return deserializeChannelStateHandle(
+                    INPUT_CHANNEL_INFO_READER, InputChannelStateHandle::new, dis, context);
+
+        } else if (MERGED_INPUT_CHANNEL_STATE_HANDLE == type) {
+            Tuple4<Integer, StreamStateHandle, Long, byte[]> t4 =
+                    deserializeMergedChannelStateHandle(dis, context);
+            return new MergedInputChannelStateHandle(t4.f0, t4.f1, t4.f2, t4.f3);
+
+        } else {
+            throw new IllegalStateException("Unknown input state handle type.");
+        }
+    }
+
+    private <Info> void serializeMergedChannelStateHandle(
+            AbstractMergedChannelStateHandle<Info, AbstractChannelStateHandle<Info>> handle,
+            DataOutputStream dos)
+            throws IOException {
+
+        dos.writeInt(handle.getSubtaskIndex());
+        dos.writeLong(handle.getStateSize());
+        serializeStreamStateHandle(handle.getDelegate(), dos);
+
+        dos.writeInt(handle.getSerializedChannelOffsets().length);
+        dos.write(handle.getSerializedChannelOffsets());
+    }
+
+    private Tuple4<Integer, StreamStateHandle, Long, byte[]> deserializeMergedChannelStateHandle(
+            DataInputStream dis, MetadataV2V3SerializerBase.DeserializationContext context)
+            throws IOException {
+
+        final int subtaskIndex = dis.readInt();
+        final long stateSize = dis.readLong();
+        final StreamStateHandle delegateHandle = deserializeStreamStateHandle(dis, context);
+
+        final int serializedChannelOffsetsLength = dis.readInt();
+        final byte[] serializedChannelOffsets = new byte[serializedChannelOffsetsLength];
+        dis.readFully(serializedChannelOffsets);
+
+        return Tuple4.of(subtaskIndex, delegateHandle, stateSize, serializedChannelOffsets);
+    }
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractMergedChannelStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractMergedChannelStateHandle.java
new file mode 100644
index 00000000000..b65cbbf84c0
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractMergedChannelStateHandle.java
@@ -0,0 +1,240 @@
+/*
+ * 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.state;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.util.Preconditions;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
+@Internal
+public abstract class AbstractMergedChannelStateHandle<
+                Info, ChannelHandle extends AbstractChannelStateHandle<Info>>
+        implements StateObject {
+
+    private static final long serialVersionUID = 1L;
+
+    /** The original subtask index before rescaling recovery. */
+    protected final int subtaskIndex;
+
+    protected final StreamStateHandle delegate;
+
+    protected final long size;
+
+    protected final byte[] serializedChannelOffsets;
+
+    private transient volatile Collection<ChannelHandle> unmergedHandles;
+
+    protected AbstractMergedChannelStateHandle(
+            int subtaskIndex,
+            StreamStateHandle delegate,
+            long size,
+            byte[] serializedChannelOffsets) {
+        this.subtaskIndex = subtaskIndex;
+        this.delegate = delegate;
+        this.size = size;
+        this.serializedChannelOffsets = serializedChannelOffsets;
+    }
+
+    protected AbstractMergedChannelStateHandle(Collection<ChannelHandle> handles) {
+        Preconditions.checkState(handles != null && handles.size() > 0);
+
+        Iterator<ChannelHandle> iterator = handles.iterator();
+        ChannelHandle handle = iterator.next();
+
+        this.subtaskIndex = handle.getSubtaskIndex();
+        this.delegate = handle.getDelegate();
+
+        long sizeAcc = 0;
+
+        ByteArrayOutputStream bos = new ByteArrayOutputStream();
+        DataOutputStream dos = new DataOutputStream(bos);
+
+        try {
+            dos.writeInt(handles.size());
+        } catch (IOException e) {
+            throw new IllegalStateException("Can not serialize Channel offsets !", e);
+        }
+
+        sizeAcc += handle.getStateSize();
+        serializeChannelOffsets(handle, dos);
+
+        while (iterator.hasNext()) {
+            handle = iterator.next();
+
+            Preconditions.checkState(
+                    this.subtaskIndex == handle.getSubtaskIndex(),
+                    "All channel state handles to be converted to merged channel state handles "
+                            + "must have the same subtask index !");
+            Preconditions.checkState(
+                    this.delegate == handle.getDelegate(),
+                    "All channel state handles to be converted to merged channel state handles "
+                            + "must have the same delegate StreamStateHandle !");
+
+            sizeAcc += handle.getStateSize();
+            serializeChannelOffsets(handle, dos);
+        }
+
+        this.size = sizeAcc;
+        this.serializedChannelOffsets = bos.toByteArray();
+        this.unmergedHandles = Collections.unmodifiableCollection(handles);
+    }
+
+    @Override
+    public void discardState() throws Exception {
+        delegate.discardState();
+    }
+
+    @Override
+    public long getStateSize() {
+        return size; // can not rely on delegate.getStateSize because it can be shared
+    }
+
+    public int getSubtaskIndex() {
+        return subtaskIndex;
+    }
+
+    public StreamStateHandle getDelegate() {
+        return delegate;
+    }
+
+    public byte[] getSerializedChannelOffsets() {
+        return serializedChannelOffsets;
+    }
+
+    public List<Info> getInfos() {
+        ensureCacheUnmergedHandles();
+        return unmergedHandles.stream().map(h -> h.getInfo()).collect(Collectors.toList());
+    }
+
+    protected abstract void writeInfo(Info info, DataOutputStream dos) throws IOException;
+
+    protected abstract Info readInfo(DataInputStream dis) throws IOException;
+
+    protected abstract ChannelHandle createUnmergedHandle(
+            int subtaskIndex,
+            StreamStateHandle delegate,
+            Info info,
+            long stateSize,
+            List<Long> offsets);
+
+    protected Collection<ChannelHandle> getUnmergedHandles() {
+        ensureCacheUnmergedHandles();
+        return unmergedHandles;
+    }
+
+    private void ensureCacheUnmergedHandles() {
+        if (unmergedHandles == null) {
+            synchronized (this) {
+                if (unmergedHandles == null) {
+                    try {
+                        restoreAndCacheUnmergedHandles();
+                    } catch (IOException e) {
+                        throw new IllegalStateException("Corrupted serializedChannelOffsets !", e);
+                    }
+                }
+            }
+        }
+    }
+
+    private void serializeChannelOffsets(ChannelHandle handle, DataOutputStream dos) {
+        try {
+            writeInfo(handle.getInfo(), dos);
+            dos.writeLong(handle.getStateSize());
+
+            dos.writeInt(handle.getOffsets().size());
+            for (Long offset : handle.getOffsets()) {
+                dos.writeLong(offset);
+            }
+        } catch (IOException e) {
+            throw new IllegalStateException("Can not serialize Channel offsets !", e);
+        }
+    }
+
+    private void restoreAndCacheUnmergedHandles() throws IOException {
+        DataInputStream dis =
+                new DataInputStream(new ByteArrayInputStream(serializedChannelOffsets));
+
+        int remainHandle = dis.readInt();
+        List<ChannelHandle> handles = new ArrayList<>(remainHandle);
+
+        while (remainHandle-- > 0) {
+            Info info = readInfo(dis);
+            long stateSize = dis.readLong();
+
+            int remainOffset = dis.readInt();
+            List<Long> offsets = new ArrayList<>(remainOffset);
+            while (remainOffset-- > 0) {
+                offsets.add(dis.readLong());
+            }
+
+            handles.add(createUnmergedHandle(subtaskIndex, delegate, info, stateSize, offsets));
+        }
+
+        unmergedHandles = Collections.unmodifiableCollection(handles);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        final AbstractMergedChannelStateHandle<?, ?> that =
+                (AbstractMergedChannelStateHandle<?, ?>) o;
+        return subtaskIndex == that.subtaskIndex
+                && delegate.equals(that.delegate)
+                && size == that.size
+                && Arrays.equals(serializedChannelOffsets, that.serializedChannelOffsets);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(
+                subtaskIndex, delegate, size, Arrays.hashCode(serializedChannelOffsets));
+    }
+
+    @Override
+    public String toString() {
+        return "AbstractMergedChannelStateHandle{"
+                + "subtaskIndex="
+                + subtaskIndex
+                + ", delegate="
+                + delegate
+                + ", serializedChannelOffsets="
+                + Arrays.toString(serializedChannelOffsets)
+                + ", size="
+                + size
+                + '}';
+    }
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/ChannelStateHelper.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/ChannelStateHelper.java
index ff72654ea31..cfbca13f2b8 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/ChannelStateHelper.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/ChannelStateHelper.java
@@ -33,6 +33,8 @@ import java.util.List;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import static java.util.stream.Collectors.groupingBy;
+
 /**
  * Utility class for channel info serialization and conversions between single channel state handle
  * and merged state handle.
@@ -72,6 +74,10 @@ public final class ChannelStateHelper {
                             if (handle instanceof AbstractChannelStateHandle) {
                                 return ((AbstractChannelStateHandle<?>) handle).getDelegate();
 
+                            } else if (handle instanceof AbstractMergedChannelStateHandle) {
+                                return ((AbstractMergedChannelStateHandle<?, ?>) handle)
+                                        .getDelegate();
+
                             } else {
                                 throw new IllegalStateException(
                                         "Not Supported state handle : " + handle.getClass());
@@ -100,6 +106,40 @@ public final class ChannelStateHelper {
                 handles.stream().map(e -> ((OutputStateHandle) e)).collect(Collectors.toList()));
     }
 
+    public static StateObjectCollection<InputStateHandle> mergeInputStateCollection(
+            Collection<InputChannelStateHandle> handles) {
+        if (handles == null || handles.isEmpty()) {
+            return StateObjectCollection.empty();
+        }
+
+        Collection<InputStateHandle> inputStateHandles =
+                handles.stream()
+                        .collect(groupingBy(AbstractChannelStateHandle::getDelegate))
+                        .values()
+                        .stream()
+                        .map(MergedInputChannelStateHandle::fromChannelHandles)
+                        .map(e -> ((InputStateHandle) e))
+                        .collect(Collectors.toSet());
+        return new StateObjectCollection<>(inputStateHandles);
+    }
+
+    public static StateObjectCollection<OutputStateHandle> mergeOutputStateCollection(
+            Collection<ResultSubpartitionStateHandle> handles) {
+        if (handles == null || handles.isEmpty()) {
+            return StateObjectCollection.empty();
+        }
+
+        Collection<OutputStateHandle> outputStateHandles =
+                handles.stream()
+                        .collect(groupingBy(AbstractChannelStateHandle::getDelegate))
+                        .values()
+                        .stream()
+                        .map(MergedResultSubpartitionStateHandle::fromChannelHandles)
+                        .map(e -> ((OutputStateHandle) e))
+                        .collect(Collectors.toSet());
+        return new StateObjectCollection<>(outputStateHandles);
+    }
+
     public static StateObjectCollection<InputChannelStateHandle> extractUnmergedInputHandles(
             OperatorSubtaskState subtaskState) {
         List<InputChannelStateHandle> inputHandles =
@@ -108,6 +148,9 @@ public final class ChannelStateHelper {
                                 h -> {
                                     if (h instanceof InputChannelStateHandle) {
                                         return Stream.of((InputChannelStateHandle) h);
+                                    } else if (h instanceof MergedInputChannelStateHandle) {
+                                        return ((MergedInputChannelStateHandle) h)
+                                                .getUnmergedHandles().stream();
                                     } else {
                                         throw new IllegalStateException(
                                                 "Invalid input channel state : " + h.getClass());
@@ -127,6 +170,9 @@ public final class ChannelStateHelper {
                                 h -> {
                                     if (h instanceof ResultSubpartitionStateHandle) {
                                         return Stream.of((ResultSubpartitionStateHandle) h);
+                                    } else if (h instanceof MergedResultSubpartitionStateHandle) {
+                                        return ((MergedResultSubpartitionStateHandle) h)
+                                                .getUnmergedHandles().stream();
                                     } else {
                                         throw new IllegalStateException(
                                                 "Invalid output channel state : " + h.getClass());
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/MergedInputChannelStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/MergedInputChannelStateHandle.java
new file mode 100644
index 00000000000..a12bee97ab4
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/MergedInputChannelStateHandle.java
@@ -0,0 +1,75 @@
+/*
+ * 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.state;
+
+import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo;
+
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.util.Collection;
+import java.util.List;
+
+import static org.apache.flink.runtime.state.ChannelStateHelper.INPUT_CHANNEL_INFO_READER;
+import static org.apache.flink.runtime.state.ChannelStateHelper.INPUT_CHANNEL_INFO_WRITER;
+
+public class MergedInputChannelStateHandle
+        extends AbstractMergedChannelStateHandle<InputChannelInfo, InputChannelStateHandle>
+        implements InputStateHandle {
+
+    private static final long serialVersionUID = 1L;
+
+    public MergedInputChannelStateHandle(
+            int subtaskIndex,
+            StreamStateHandle delegate,
+            long size,
+            byte[] serializedChannelOffsets) {
+        super(subtaskIndex, delegate, size, serializedChannelOffsets);
+    }
+
+    @Override
+    protected void writeInfo(InputChannelInfo inputChannelInfo, DataOutputStream dos)
+            throws IOException {
+        INPUT_CHANNEL_INFO_WRITER.accept(inputChannelInfo, dos);
+    }
+
+    @Override
+    protected InputChannelInfo readInfo(DataInputStream dis) throws IOException {
+        return INPUT_CHANNEL_INFO_READER.apply(dis);
+    }
+
+    @Override
+    protected InputChannelStateHandle createUnmergedHandle(
+            int subtaskIndex,
+            StreamStateHandle delegate,
+            InputChannelInfo info,
+            long stateSize,
+            List<Long> offsets) {
+        return new InputChannelStateHandle(subtaskIndex, info, delegate, offsets, stateSize);
+    }
+
+    private MergedInputChannelStateHandle(Collection<InputChannelStateHandle> handles) {
+        super(handles);
+    }
+
+    public static MergedInputChannelStateHandle fromChannelHandles(
+            Collection<InputChannelStateHandle> handles) {
+        return new MergedInputChannelStateHandle(handles);
+    }
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/MergedResultSubpartitionStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/MergedResultSubpartitionStateHandle.java
new file mode 100644
index 00000000000..8d2a335e6f8
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/MergedResultSubpartitionStateHandle.java
@@ -0,0 +1,76 @@
+/*
+ * 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.state;
+
+import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo;
+
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.util.Collection;
+import java.util.List;
+
+import static org.apache.flink.runtime.state.ChannelStateHelper.RESULT_SUBPARTITION_INFO_READER;
+import static org.apache.flink.runtime.state.ChannelStateHelper.RESULT_SUBPARTITION_INFO_WRITER;
+
+public class MergedResultSubpartitionStateHandle
+        extends AbstractMergedChannelStateHandle<
+                ResultSubpartitionInfo, ResultSubpartitionStateHandle>
+        implements OutputStateHandle {
+
+    private static final long serialVersionUID = 1L;
+
+    public MergedResultSubpartitionStateHandle(
+            int subtaskIndex,
+            StreamStateHandle delegate,
+            long size,
+            byte[] serializedChannelOffsets) {
+        super(subtaskIndex, delegate, size, serializedChannelOffsets);
+    }
+
+    @Override
+    protected void writeInfo(ResultSubpartitionInfo resultSubpartitionInfo, DataOutputStream dos)
+            throws IOException {
+        RESULT_SUBPARTITION_INFO_WRITER.accept(resultSubpartitionInfo, dos);
+    }
+
+    @Override
+    protected ResultSubpartitionInfo readInfo(DataInputStream dis) throws IOException {
+        return RESULT_SUBPARTITION_INFO_READER.apply(dis);
+    }
+
+    @Override
+    protected ResultSubpartitionStateHandle createUnmergedHandle(
+            int subtaskIndex,
+            StreamStateHandle delegate,
+            ResultSubpartitionInfo info,
+            long stateSize,
+            List<Long> offsets) {
+        return new ResultSubpartitionStateHandle(subtaskIndex, info, delegate, offsets, stateSize);
+    }
+
+    private MergedResultSubpartitionStateHandle(Collection<ResultSubpartitionStateHandle> handles) {
+        super(handles);
+    }
+
+    public static MergedResultSubpartitionStateHandle fromChannelHandles(
+            Collection<ResultSubpartitionStateHandle> handles) {
+        return new MergedResultSubpartitionStateHandle(handles);
+    }
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV1Test.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV1Test.java
index 5823cf5beaa..e7c31768cd0 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV1Test.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV1Test.java
@@ -18,94 +18,26 @@
 
 package org.apache.flink.runtime.checkpoint.metadata;
 
-import org.apache.flink.api.java.tuple.Tuple5;
-import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo;
-import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo;
-import org.apache.flink.runtime.state.InputChannelStateHandle;
-import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
-import org.apache.flink.runtime.state.StreamStateHandle;
-import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
-
 import org.junit.jupiter.api.Test;
 
-import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
-import java.io.DataInputStream;
-import java.io.DataOutputStream;
 import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.concurrent.ThreadLocalRandom;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.apache.flink.runtime.checkpoint.metadata.ChannelStateTestUtils.testSerializeInputHandle;
+import static org.apache.flink.runtime.checkpoint.metadata.ChannelStateTestUtils.testSerializeOutputHandle;
 
 class ChannelStateHandleSerializerV1Test {
 
     @Test
-    void testSerializeInputHandle() throws IOException {
-        ChannelStateHandleSerializerV1 serializer = new ChannelStateHandleSerializerV1();
-
-        Tuple5<Integer, Integer, Integer, StreamStateHandle, List<Long>> handleFields =
-                generateRandomChannelHandleFields();
-
-        InputChannelStateHandle inputHandle =
-                new InputChannelStateHandle(
-                        handleFields.f0,
-                        new InputChannelInfo(handleFields.f1, handleFields.f2),
-                        handleFields.f3,
-                        handleFields.f4,
-                        handleFields.f3.getStateSize());
-
-        ByteArrayOutputStream inMemBuffer = new ByteArrayOutputStream();
-        DataOutputStream outputStream = new DataOutputStream(inMemBuffer);
-        serializer.serialize(inputHandle, outputStream);
-
-        DataInputStream inputStream =
-                new DataInputStream(new ByteArrayInputStream(inMemBuffer.toByteArray()));
-        assertEquals(inputHandle, serializer.deserializeInputStateHandle(inputStream, null));
+    void testWithInputChannelStateHandle() throws IOException {
+        testSerializeInputHandle(
+                new ChannelStateHandleSerializerV1(),
+                ChannelStateTestUtils::randomInputChannelStateHandle);
     }
 
     @Test
-    void testSerializeOutputHandle() throws IOException {
-        ChannelStateHandleSerializerV1 serializer = new ChannelStateHandleSerializerV1();
-
-        Tuple5<Integer, Integer, Integer, StreamStateHandle, List<Long>> handleFields =
-                generateRandomChannelHandleFields();
-
-        ResultSubpartitionStateHandle outputHandle =
-                new ResultSubpartitionStateHandle(
-                        handleFields.f0,
-                        new ResultSubpartitionInfo(handleFields.f1, handleFields.f2),
-                        handleFields.f3,
-                        handleFields.f4,
-                        handleFields.f3.getStateSize());
-
-        ByteArrayOutputStream inMemBuffer = new ByteArrayOutputStream();
-        DataOutputStream outputStream = new DataOutputStream(inMemBuffer);
-        serializer.serialize(outputHandle, outputStream);
-
-        DataInputStream inputStream =
-                new DataInputStream(new ByteArrayInputStream(inMemBuffer.toByteArray()));
-        assertEquals(outputHandle, serializer.deserializeInputStateHandle(inputStream, null));
-    }
-
-    private Tuple5<Integer, Integer, Integer, StreamStateHandle, List<Long>>
-            generateRandomChannelHandleFields() {
-        ThreadLocalRandom random = ThreadLocalRandom.current();
-
-        int subtaskIndex = random.nextInt();
-        int infoIdx1 = random.nextInt();
-        int infoIdx2 = random.nextInt();
-
-        byte[] randomBytes = String.valueOf(random.nextLong()).getBytes();
-        StreamStateHandle delegate = new ByteStreamStateHandle("delegated", randomBytes);
-
-        int offsetSize = random.nextInt(10, 20);
-        List<Long> offsets = new ArrayList<>(offsetSize);
-        for (int i = 0; i < offsetSize; i++) {
-            offsets.add(random.nextLong());
-        }
-
-        return Tuple5.of(subtaskIndex, infoIdx1, infoIdx2, delegate, offsets);
+    void testWithResultSubpartitionStateHandle() throws IOException {
+        testSerializeOutputHandle(
+                new ChannelStateHandleSerializerV1(),
+                ChannelStateTestUtils::randomResultSubpartitionStateHandle);
     }
 }
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV2Test.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV2Test.java
new file mode 100644
index 00000000000..80f78a368bf
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateHandleSerializerV2Test.java
@@ -0,0 +1,57 @@
+/*
+ * 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.metadata;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+
+import static org.apache.flink.runtime.checkpoint.metadata.ChannelStateTestUtils.testSerializeInputHandle;
+import static org.apache.flink.runtime.checkpoint.metadata.ChannelStateTestUtils.testSerializeOutputHandle;
+
+class ChannelStateHandleSerializerV2Test {
+
+    @Test
+    void testWithInputChannelStateHandle() throws IOException {
+        testSerializeInputHandle(
+                new ChannelStateHandleSerializerV2(),
+                ChannelStateTestUtils::randomInputChannelStateHandle);
+    }
+
+    @Test
+    void testWithResultSubpartitionStateHandle() throws IOException {
+        testSerializeOutputHandle(
+                new ChannelStateHandleSerializerV2(),
+                ChannelStateTestUtils::randomResultSubpartitionStateHandle);
+    }
+
+    @Test
+    void testWithMergedInputChannelStateHandle() throws IOException {
+        testSerializeInputHandle(
+                new ChannelStateHandleSerializerV2(),
+                ChannelStateTestUtils::randomMergedInputChannelStateHandle);
+    }
+
+    @Test
+    void testWithMergedResultSubpartitionStateHandle() throws IOException {
+        testSerializeOutputHandle(
+                new ChannelStateHandleSerializerV2(),
+                ChannelStateTestUtils::randomMergedResultSubpartitionStateHandle);
+    }
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateTestUtils.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateTestUtils.java
new file mode 100644
index 00000000000..1a8ad9bee82
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/metadata/ChannelStateTestUtils.java
@@ -0,0 +1,175 @@
+/*
+ * 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.metadata;
+
+import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo;
+import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo;
+import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.InputStateHandle;
+import org.apache.flink.runtime.state.MergedInputChannelStateHandle;
+import org.apache.flink.runtime.state.MergedResultSubpartitionStateHandle;
+import org.apache.flink.runtime.state.OutputStateHandle;
+import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.Supplier;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+/** A collection of utility methods for testing about channel state. */
+public class ChannelStateTestUtils {
+    private static final int SUBTASK_INDEX_UPPER = 10000;
+
+    public static InputChannelStateHandle randomInputChannelStateHandle() {
+        ThreadLocalRandom random = ThreadLocalRandom.current();
+        StreamStateHandle delegated = randomDelegatedHandle();
+
+        return new InputChannelStateHandle(
+                random.nextInt(0, SUBTASK_INDEX_UPPER),
+                new InputChannelInfo(
+                        random.nextInt(0, SUBTASK_INDEX_UPPER),
+                        random.nextInt(0, SUBTASK_INDEX_UPPER)),
+                delegated,
+                randomOffsets(),
+                delegated.getStateSize());
+    }
+
+    public static ResultSubpartitionStateHandle randomResultSubpartitionStateHandle() {
+        ThreadLocalRandom random = ThreadLocalRandom.current();
+        StreamStateHandle delegated = randomDelegatedHandle();
+
+        return new ResultSubpartitionStateHandle(
+                random.nextInt(0, SUBTASK_INDEX_UPPER),
+                new ResultSubpartitionInfo(
+                        random.nextInt(0, SUBTASK_INDEX_UPPER),
+                        random.nextInt(0, SUBTASK_INDEX_UPPER)),
+                delegated,
+                randomOffsets(),
+                delegated.getStateSize());
+    }
+
+    public static List<InputChannelStateHandle> randomInputChannelStateHandlesFromSameSubtask() {
+        ThreadLocalRandom random = ThreadLocalRandom.current();
+
+        int subtaskIndex = random.nextInt(0, SUBTASK_INDEX_UPPER);
+        StreamStateHandle delegated = randomDelegatedHandle();
+
+        int channelHandleCount = random.nextInt(20, 50);
+        List<InputChannelStateHandle> inputHandles = new ArrayList<>(channelHandleCount);
+        for (int i = 0; i < channelHandleCount; i++) {
+            inputHandles.add(
+                    new InputChannelStateHandle(
+                            subtaskIndex,
+                            new InputChannelInfo(
+                                    random.nextInt(0, SUBTASK_INDEX_UPPER),
+                                    random.nextInt(0, SUBTASK_INDEX_UPPER)),
+                            delegated,
+                            randomOffsets(),
+                            delegated.getStateSize()));
+        }
+
+        return inputHandles;
+    }
+
+    public static MergedInputChannelStateHandle randomMergedInputChannelStateHandle() {
+        return MergedInputChannelStateHandle.fromChannelHandles(
+                randomInputChannelStateHandlesFromSameSubtask());
+    }
+
+    public static List<ResultSubpartitionStateHandle>
+            randomResultSubpartitionStateHandlesFromSameSubtask() {
+        ThreadLocalRandom random = ThreadLocalRandom.current();
+
+        int subtaskIndex = random.nextInt(0, SUBTASK_INDEX_UPPER);
+        StreamStateHandle delegated = randomDelegatedHandle();
+
+        int channelHandleCount = random.nextInt(20, 50);
+        List<ResultSubpartitionStateHandle> outputHandles = new ArrayList<>(channelHandleCount);
+        for (int i = 0; i < channelHandleCount; i++) {
+            outputHandles.add(
+                    new ResultSubpartitionStateHandle(
+                            subtaskIndex,
+                            new ResultSubpartitionInfo(
+                                    random.nextInt(0, SUBTASK_INDEX_UPPER),
+                                    random.nextInt(0, SUBTASK_INDEX_UPPER)),
+                            delegated,
+                            randomOffsets(),
+                            delegated.getStateSize()));
+        }
+
+        return outputHandles;
+    }
+
+    public static MergedResultSubpartitionStateHandle randomMergedResultSubpartitionStateHandle() {
+        return MergedResultSubpartitionStateHandle.fromChannelHandles(
+                randomResultSubpartitionStateHandlesFromSameSubtask());
+    }
+
+    private static StreamStateHandle randomDelegatedHandle() {
+        byte[] randomBytes = String.valueOf(ThreadLocalRandom.current().nextLong()).getBytes();
+        return new ByteStreamStateHandle("delegated", randomBytes);
+    }
+
+    private static List<Long> randomOffsets() {
+        ThreadLocalRandom random = ThreadLocalRandom.current();
+        int offsetSize = random.nextInt(10, 20);
+        List<Long> offsets = new ArrayList<>(offsetSize);
+        for (int i = 0; i < offsetSize; i++) {
+            offsets.add(random.nextLong());
+        }
+        return offsets;
+    }
+
+    static void testSerializeInputHandle(
+            ChannelStateHandleSerializer serializer, Supplier<InputStateHandle> handleGenerator)
+            throws IOException {
+        InputStateHandle inputHandle = handleGenerator.get();
+
+        ByteArrayOutputStream inMemBuffer = new ByteArrayOutputStream();
+        DataOutputStream outputStream = new DataOutputStream(inMemBuffer);
+        serializer.serialize(inputHandle, outputStream);
+
+        DataInputStream inputStream =
+                new DataInputStream(new ByteArrayInputStream(inMemBuffer.toByteArray()));
+        assertEquals(inputHandle, serializer.deserializeInputStateHandle(inputStream, null));
+    }
+
+    static void testSerializeOutputHandle(
+            ChannelStateHandleSerializer serializer, Supplier<OutputStateHandle> handleGenerator)
+            throws IOException {
+        OutputStateHandle outputHandle = handleGenerator.get();
+
+        ByteArrayOutputStream inMemBuffer = new ByteArrayOutputStream();
+        DataOutputStream outputStream = new DataOutputStream(inMemBuffer);
+        serializer.serialize(outputHandle, outputStream);
+
+        DataInputStream inputStream =
+                new DataInputStream(new ByteArrayInputStream(inMemBuffer.toByteArray()));
+        assertEquals(outputHandle, serializer.deserializeOutputStateHandle(inputStream, null));
+    }
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/MergedChannelStateHandleTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/MergedChannelStateHandleTest.java
new file mode 100644
index 00000000000..8f62b2c2350
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/MergedChannelStateHandleTest.java
@@ -0,0 +1,49 @@
+/*
+ * 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.state;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.apache.flink.runtime.checkpoint.metadata.ChannelStateTestUtils.randomInputChannelStateHandlesFromSameSubtask;
+import static org.apache.flink.runtime.checkpoint.metadata.ChannelStateTestUtils.randomResultSubpartitionStateHandlesFromSameSubtask;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+class MergedChannelStateHandleTest {
+
+    @Test
+    void testMergeInputChannelStateHandles() {
+        List<InputChannelStateHandle> inputHandles =
+                randomInputChannelStateHandlesFromSameSubtask();
+        MergedInputChannelStateHandle mergedHandle =
+                MergedInputChannelStateHandle.fromChannelHandles(inputHandles);
+        assertEquals(inputHandles, new ArrayList<>(mergedHandle.getUnmergedHandles()));
+    }
+
+    @Test
+    void testMergeResultSubpartitionStateHandles() {
+        List<ResultSubpartitionStateHandle> outputHandles =
+                randomResultSubpartitionStateHandlesFromSameSubtask();
+        MergedResultSubpartitionStateHandle mergedHandle =
+                MergedResultSubpartitionStateHandle.fromChannelHandles(outputHandles);
+        assertEquals(outputHandles, new ArrayList<>(mergedHandle.getUnmergedHandles()));
+    }
+}