From 4fddae395d7b6e9493209af0f7edc6a5fa55bec0 Mon Sep 17 00:00:00 2001 From: Zakelly Date: Sat, 29 Aug 2026 21:27:01 +0800 Subject: [PATCH 1/2] [FLINK-39977][runtime] Recovery of merged channel state handles --- ...askFileMergingManagerRestoreOperation.java | 11 ++- .../FileMergingSnapshotManagerTestBase.java | 67 ++++++++++++++++++- 2 files changed, 74 insertions(+), 4 deletions(-) diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/filemerging/SubtaskFileMergingManagerRestoreOperation.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/filemerging/SubtaskFileMergingManagerRestoreOperation.java index 36c064efb84ac6..9dd8c59d02893b 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/filemerging/SubtaskFileMergingManagerRestoreOperation.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/filemerging/SubtaskFileMergingManagerRestoreOperation.java @@ -35,6 +35,8 @@ import java.util.function.Function; import java.util.stream.Stream; +import static org.apache.flink.runtime.state.ChannelStateHelper.collectUniqueDisposableInChannelState; + /** * Restore operation that restores file-merging information belonging to one subtask for {@link * FileMergingSnapshotManager}. @@ -89,10 +91,15 @@ public void restore() { subtaskState.getRawOperatorState().stream()) .flatMap(this::getChildrenStreamHandles); - // TODO support channel state restore for unaligned checkpoint. + Stream channelStateHandles = + collectUniqueDisposableInChannelState( + Stream.of( + subtaskState.getInputChannelState(), + subtaskState.getUpstreamOutputBufferState(), + subtaskState.getResultSubpartitionState())); Stream segmentStateHandles = - Stream.of(keyedStateHandles, operatorStateHandles) + Stream.of(keyedStateHandles, operatorStateHandles, channelStateHandles) .flatMap(Function.identity()) .filter( handle -> diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/filemerging/FileMergingSnapshotManagerTestBase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/filemerging/FileMergingSnapshotManagerTestBase.java index 185300ff64169c..8abb223e6999ec 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/filemerging/FileMergingSnapshotManagerTestBase.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/filemerging/FileMergingSnapshotManagerTestBase.java @@ -27,7 +27,10 @@ import org.apache.flink.core.fs.Path; import org.apache.flink.core.fs.local.LocalFileSystem; 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.InputChannelInfo; +import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo; import org.apache.flink.runtime.checkpoint.filemerging.FileMergingSnapshotManager.SpaceStat; import org.apache.flink.runtime.checkpoint.filemerging.FileMergingSnapshotManager.SubtaskKey; import org.apache.flink.runtime.clusterframework.types.ResourceID; @@ -37,10 +40,14 @@ import org.apache.flink.runtime.state.CheckpointedStateScope; import org.apache.flink.runtime.state.IncrementalKeyedStateHandle; import org.apache.flink.runtime.state.IncrementalRemoteKeyedStateHandle; +import org.apache.flink.runtime.state.InputChannelStateHandle; import org.apache.flink.runtime.state.KeyGroupRange; import org.apache.flink.runtime.state.KeyGroupRangeOffsets; import org.apache.flink.runtime.state.KeyGroupsStateHandle; +import org.apache.flink.runtime.state.MergedInputChannelStateHandle; 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.filemerging.FileMergingOperatorStreamStateHandle; import org.apache.flink.runtime.state.filemerging.SegmentFileStateHandle; import org.apache.flink.runtime.state.filesystem.FileMergingCheckpointStateOutputStream; @@ -51,6 +58,7 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -539,8 +547,8 @@ public void testRestore() throws Exception { assertThat(stateFiles.size()).isEqualTo(1); Set restoreFileSet = stateFiles.get(checkpointId); assertThat(restoreFileSet).isNotNull(); - // 2 operators * (2 keyed state + 2 operator state) - assertThat(restoreFileSet.size()).isEqualTo(8); + // 2 operators * (2 keyed state + 2 operator state + 3 channel state delegates) + assertThat(restoreFileSet.size()).isEqualTo(14); assertThat(fmsm.spaceStat).isEqualTo(oldSpaceStat); for (LogicalFile file : restoreFileSet) { assertThat(fmsm.getLogicalFile(file.getFileId())).isEqualTo(file); @@ -741,11 +749,66 @@ private OperatorSubtaskState buildOperatorSubtaskState( CheckpointedStateScope.EXCLUSIVE, closeableRegistry)); + SegmentFileStateHandle inputChannelStateHandle = + buildOneSegmentFileHandle( + checkpointId, fmsm, CheckpointedStateScope.EXCLUSIVE, closeableRegistry); + InputChannelStateHandle inputChannel1 = + new InputChannelStateHandle( + 0, + new InputChannelInfo(0, 0), + inputChannelStateHandle, + Collections.singletonList(0L), + 16L); + InputChannelStateHandle inputChannel2 = + new InputChannelStateHandle( + 0, + new InputChannelInfo(0, 1), + inputChannelStateHandle, + Collections.singletonList(16L), + 16L); + + SegmentFileStateHandle upstreamOutputBufferStateHandle = + buildOneSegmentFileHandle( + checkpointId, fmsm, CheckpointedStateScope.EXCLUSIVE, closeableRegistry); + InputChannelStateHandle upstreamOutputBuffer = + new InputChannelStateHandle( + 0, + new InputChannelInfo(1, 0), + upstreamOutputBufferStateHandle, + Collections.singletonList(0L), + upstreamOutputBufferStateHandle.getStateSize()); + + SegmentFileStateHandle resultSubpartitionStateHandle = + buildOneSegmentFileHandle( + checkpointId, fmsm, CheckpointedStateScope.EXCLUSIVE, closeableRegistry); + ResultSubpartitionStateHandle resultSubpartition1 = + new ResultSubpartitionStateHandle( + 0, + new ResultSubpartitionInfo(0, 0), + resultSubpartitionStateHandle, + Collections.singletonList(0L), + 16L); + ResultSubpartitionStateHandle resultSubpartition2 = + new ResultSubpartitionStateHandle( + 0, + new ResultSubpartitionInfo(0, 1), + resultSubpartitionStateHandle, + Collections.singletonList(16L), + 16L); + return OperatorSubtaskState.builder() .setManagedKeyedState(keyedStateHandle1) .setRawKeyedState(keyedStateHandle2) .setManagedOperatorState(operatorStateHandle1) .setRawOperatorState(operatorStateHandle2) + .setInputChannelState( + StateObjectCollection.singleton( + MergedInputChannelStateHandle.fromChannelHandles( + Arrays.asList(inputChannel1, inputChannel2)))) + .setUpstreamOutputBufferState(StateObjectCollection.singleton(upstreamOutputBuffer)) + .setResultSubpartitionState( + new StateObjectCollection( + Arrays.asList(resultSubpartition1, resultSubpartition2))) .build(); } From 38ebba0f11795419962903d3cf910c37e021dab2 Mon Sep 17 00:00:00 2001 From: Zakelly Date: Sat, 29 Aug 2026 22:08:25 +0800 Subject: [PATCH 2/2] [FLINK-39977][runtime] Add ITCase for file-merged channel state recovery Generated-by: OpenAI Codex (GPT-5.6 Sol) --- .../FileMergingChannelStateITCase.java | 291 ++++++++++++++++++ 1 file changed, 291 insertions(+) create mode 100644 flink-tests/src/test/java/org/apache/flink/test/checkpointing/FileMergingChannelStateITCase.java diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/FileMergingChannelStateITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/FileMergingChannelStateITCase.java new file mode 100644 index 00000000000000..bcb90bde752a54 --- /dev/null +++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/FileMergingChannelStateITCase.java @@ -0,0 +1,291 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.test.checkpointing; + +import org.apache.flink.api.common.JobExecutionResult; +import org.apache.flink.api.common.functions.RichMapFunction; +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDescriptor; +import org.apache.flink.api.common.typeutils.base.array.LongPrimitiveArraySerializer; +import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.configuration.CheckpointingOptions; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.core.execution.CheckpointingMode; +import org.apache.flink.core.execution.JobClient; +import org.apache.flink.runtime.checkpoint.OperatorState; +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState; +import org.apache.flink.runtime.checkpoint.metadata.CheckpointMetadata; +import org.apache.flink.runtime.minicluster.MiniCluster; +import org.apache.flink.runtime.state.FunctionInitializationContext; +import org.apache.flink.runtime.state.FunctionSnapshotContext; +import org.apache.flink.runtime.state.StreamStateHandle; +import org.apache.flink.runtime.state.filemerging.SegmentFileStateHandle; +import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration; +import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.api.functions.sink.legacy.SinkFunction; +import org.apache.flink.streaming.util.CheckpointStorageUtils; +import org.apache.flink.streaming.util.RestartStrategyUtils; +import org.apache.flink.test.junit5.InjectMiniCluster; +import org.apache.flink.test.junit5.MiniClusterExtension; +import org.apache.flink.test.util.TestUtils; +import org.apache.flink.testutils.junit.SharedObjectsExtension; +import org.apache.flink.testutils.junit.SharedReference; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.api.io.TempDir; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicLongArray; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +import static org.apache.flink.runtime.state.ChannelStateHelper.collectUniqueDisposableInChannelState; +import static org.apache.flink.runtime.testutils.CommonTestUtils.waitUntilCondition; +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests recovery of file-merged channel state after the TaskManager is replaced. */ +class FileMergingChannelStateITCase { + + private static final int TASK_MANAGER_COUNT = 3; + private static final int WORD_COUNT = 16; + private static final long RECORD_COUNT = 16_000L; + private static final long EXPECTED_COUNT_PER_WORD = RECORD_COUNT / WORD_COUNT; + private static final String SLOW_MAPPER_UID = "slow-word-mapper"; + + private static final List WORDS = + Arrays.asList( + "apple", + "banana", + "cherry", + "date", + "elderberry", + "fig", + "grape", + "honeydew", + "kiwi", + "lemon", + "mango", + "nectarine", + "orange", + "papaya", + "quince", + "raspberry"); + + @TempDir private java.nio.file.Path checkpointDirectory; + + @RegisterExtension + private final SharedObjectsExtension sharedObjects = SharedObjectsExtension.create(); + + @RegisterExtension + private static final MiniClusterExtension MINI_CLUSTER_EXTENSION = + new MiniClusterExtension( + new MiniClusterResourceConfiguration.Builder() + .setConfiguration(clusterConfiguration()) + .setNumberTaskManagers(TASK_MANAGER_COUNT) + .setNumberSlotsPerTaskManager(1) + .build()); + + private static Configuration clusterConfiguration() { + return new Configuration() + .set(CheckpointingOptions.FILE_MERGING_ENABLED, true) + .set(CheckpointingOptions.FILE_MERGING_ACROSS_BOUNDARY, false) + .set(CheckpointingOptions.MAX_RETAINED_CHECKPOINTS, 1); + } + + @Test + void testRestoreFileMergedChannelState(@InjectMiniCluster MiniCluster miniCluster) + throws Exception { + final SharedReference mappedRecords = sharedObjects.add(new AtomicLong()); + final SharedReference result = + sharedObjects.add(new AtomicLongArray(WORD_COUNT)); + final StreamExecutionEnvironment env = createEnvironment(mappedRecords, result); + final JobClient jobClient = env.executeAsync("file-merging-channel-state-word-count"); + + try { + waitUntilCondition(() -> mappedRecords.get().get() >= 100L, 100L, 300); + + final String checkpointPath = + miniCluster.triggerCheckpoint(jobClient.getJobID()).get(30, TimeUnit.SECONDS); + assertFileMergedChannelState(TestUtils.loadCheckpointMetadata(checkpointPath)); + + for (int i = 0; i < TASK_MANAGER_COUNT; i++) { + miniCluster.terminateTaskManager(0).get(30, TimeUnit.SECONDS); + } + for (int i = 0; i < TASK_MANAGER_COUNT; i++) { + miniCluster.startTaskManager(); + } + + final long mappedRecordsBeforeRecovery = mappedRecords.get().get(); + waitUntilCondition( + () -> mappedRecords.get().get() >= mappedRecordsBeforeRecovery + 1_000L, + 100L, + 300); + final String postRecoveryCheckpointPath = + miniCluster.triggerCheckpoint(jobClient.getJobID()).get(30, TimeUnit.SECONDS); + assertFileMergedChannelState( + TestUtils.loadCheckpointMetadata(postRecoveryCheckpointPath)); + + final JobExecutionResult executionResult = + jobClient.getJobExecutionResult().get(2, TimeUnit.MINUTES); + assertThat(executionResult.getJobID()).isEqualTo(jobClient.getJobID()); + for (int i = 0; i < WORD_COUNT; i++) { + assertThat(result.get().get(i)) + .as("final count for word %s", WORDS.get(i)) + .isEqualTo(EXPECTED_COUNT_PER_WORD); + } + } finally { + if (!jobClient.getJobExecutionResult().isDone()) { + jobClient.cancel().get(30, TimeUnit.SECONDS); + } + } + } + + private StreamExecutionEnvironment createEnvironment( + SharedReference mappedRecords, SharedReference result) { + final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(1); + env.enableCheckpointing(TimeUnit.DAYS.toMillis(1), CheckpointingMode.EXACTLY_ONCE); + env.getCheckpointConfig().enableUnalignedCheckpoints(); + env.getCheckpointConfig().setAlignedCheckpointTimeout(Duration.ZERO); + CheckpointStorageUtils.configureFileSystemCheckpointStorage( + env, checkpointDirectory.toUri()); + RestartStrategyUtils.configureFixedDelayRestartStrategy(env, 3, 100L); + + env.fromSequence(0L, RECORD_COUNT - 1L) + .setParallelism(1) + .slotSharingGroup("source") + .rebalance() + .map(new SlowWordMapper(mappedRecords)) + .setParallelism(1) + .slotSharingGroup("channel") + .uid(SLOW_MAPPER_UID) + .keyBy(value -> value.f0) + .sum(1) + .setParallelism(1) + .slotSharingGroup("state") + .addSink(new ResultSink(result)) + .setParallelism(1) + .slotSharingGroup("state"); + return env; + } + + private static void assertFileMergedChannelState(CheckpointMetadata metadata) { + final List channelStateDelegates = new ArrayList<>(); + final List slowMapperChannelStateDelegates = new ArrayList<>(); + for (OperatorState operatorState : metadata.getOperatorStates()) { + for (OperatorSubtaskState subtaskState : operatorState.getStates()) { + final List subtaskChannelStateDelegates = + collectUniqueDisposableInChannelState( + Stream.of( + subtaskState.getInputChannelState(), + subtaskState.getUpstreamOutputBufferState(), + subtaskState.getResultSubpartitionState())) + .collect(Collectors.toList()); + channelStateDelegates.addAll(subtaskChannelStateDelegates); + if (operatorState.getOperatorUid().filter(SLOW_MAPPER_UID::equals).isPresent()) { + collectUniqueDisposableInChannelState( + Stream.of(subtaskState.getInputChannelState())) + .forEach(slowMapperChannelStateDelegates::add); + } + } + } + + assertThat(channelStateDelegates) + .as("channel state delegates in the checkpoint") + .isNotEmpty() + .allSatisfy( + handle -> assertThat(handle).isInstanceOf(SegmentFileStateHandle.class)); + assertThat(channelStateDelegates.stream().mapToLong(StreamStateHandle::getStateSize).sum()) + .isPositive(); + assertThat(slowMapperChannelStateDelegates) + .as("channel state delegates belonging to the stateless slow mapper") + .isNotEmpty(); + } + + private static final class SlowWordMapper extends RichMapFunction> { + + private static final long serialVersionUID = 1L; + + private final SharedReference mappedRecords; + + private SlowWordMapper(SharedReference mappedRecords) { + this.mappedRecords = mappedRecords; + } + + @Override + public Tuple2 map(Long value) throws Exception { + Thread.sleep(1L); + mappedRecords.get().incrementAndGet(); + return Tuple2.of(WORDS.get((int) (value % WORD_COUNT)), 1L); + } + } + + private static final class ResultSink + implements SinkFunction>, CheckpointedFunction { + + private static final long serialVersionUID = 1L; + + private final SharedReference result; + + private transient ListState resultState; + private long[] counts = new long[WORD_COUNT]; + + private ResultSink(SharedReference result) { + this.result = result; + } + + @Override + public void invoke(Tuple2 value, Context context) { + final int wordIndex = WORDS.indexOf(value.f0); + counts[wordIndex] = value.f1; + result.get().set(wordIndex, value.f1); + } + + @Override + public void snapshotState(FunctionSnapshotContext context) throws Exception { + resultState.update(Arrays.asList(counts.clone())); + } + + @Override + public void initializeState(FunctionInitializationContext context) throws Exception { + resultState = + context.getOperatorStateStore() + .getListState( + new ListStateDescriptor<>( + "word-count-results", + LongPrimitiveArraySerializer.INSTANCE)); + counts = new long[WORD_COUNT]; + if (context.isRestored()) { + for (long[] restoredCounts : resultState.get()) { + counts = restoredCounts.clone(); + } + } + for (int i = 0; i < WORD_COUNT; i++) { + result.get().set(i, counts[i]); + } + } + } +}