This is an automated email from the ASF dual-hosted git repository. 1996fanrui pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/flink.git
commit ddeb17e52ce13cecdc3e486971750e1807b7cbd1 Author: Rui Fan <[email protected]> AuthorDate: Mon Jul 6 02:14:11 2026 +0200 [FLINK-40080][checkpoint] ChannelStateWriter#addInputDataFromSpill: replay spilled segments into the checkpoint ChannelStateWriter gains addInputDataFromSpill(long checkpointId, FetchedChannelStateReader reader); the NO_OP implementation closes the reader. ChannelStateWriterImpl enqueues ChannelStateWriteRequest.replayInputDataFromSpill (a CheckpointInProgressRequest whose cancel action closes the reader). ChannelStateCheckpointWriter.writeInputFromSpill loops nextSegment(), writes each bodyStream() via the new ChannelStateSerializer.writeData(DataOutputStream, InputStream, int) overload (length prefix + transferTo, fail-loud on short read), records offset/size into pendingResult.getInputChannelOffsets() demuxed by channelInfo, and closes the reader in a finally block. Failures propagate via ChannelStateWriteResult. --- .../channel/ChannelStateCheckpointWriter.java | 47 ++++++++ .../checkpoint/channel/ChannelStateSerializer.java | 14 +++ .../channel/ChannelStateWriteRequest.java | 14 +++ .../checkpoint/channel/ChannelStateWriter.java | 20 ++++ .../checkpoint/channel/ChannelStateWriterImpl.java | 7 ++ .../runtime/io/checkpointing/ChannelState.java | 18 ++- .../io/checkpointing/InputProcessorUtil.java | 9 +- .../SingleCheckpointBarrierHandler.java | 9 +- .../channel/ChannelStateWriterImplTest.java | 119 ++++++++++++++++++++ .../checkpoint/channel/MockChannelStateWriter.java | 10 ++ .../runtime/io/checkpointing/ChannelStateTest.java | 124 +++++++++++++++++++-- .../checkpointing/TestBarrierHandlerFactory.java | 1 + 12 files changed, 371 insertions(+), 21 deletions(-) diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateCheckpointWriter.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateCheckpointWriter.java index 4173bb7140e..a58e924fbfc 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateCheckpointWriter.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateCheckpointWriter.java @@ -19,6 +19,7 @@ package org.apache.flink.runtime.checkpoint.channel; import org.apache.flink.annotation.VisibleForTesting; import org.apache.flink.runtime.checkpoint.CheckpointException; +import org.apache.flink.runtime.checkpoint.channel.FetchedChannelStateReader.SpillSegment; import org.apache.flink.runtime.io.network.buffer.Buffer; import org.apache.flink.runtime.io.network.logger.NetworkActionsLogger; import org.apache.flink.runtime.jobgraph.JobVertexID; @@ -41,6 +42,7 @@ import java.util.HashMap; import java.util.HashSet; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.Set; import static org.apache.flink.runtime.checkpoint.CheckpointFailureReason.CHANNEL_STATE_SHARED_STREAM_EXCEPTION; @@ -161,6 +163,51 @@ class ChannelStateCheckpointWriter { } } + void writeInputFromSpill( + JobVertexID jobVertexID, int subtaskIndex, FetchedChannelStateReader reader) { + try { + if (isDone()) { + return; + } + ChannelStatePendingResult pendingResult = + getChannelStatePendingResult(jobVertexID, subtaskIndex); + runWithChecks( + () -> { + checkState(!pendingResult.isAllInputsReceived()); + String action = "ChannelStateCheckpointWriter#writeInputFromSpill"; + Optional<SpillSegment> next; + while ((next = reader.advanceAndGetNextSegment()).isPresent()) { + SpillSegment seg = next.get(); + long offset = checkpointStream.getPos(); + try (AutoCloseable ignored = + NetworkActionsLogger.measureIO(action, seg.channelInfo())) { + serializer.writeData(dataStream, seg.bodyStream(), seg.length()); + } + long size = checkpointStream.getPos() - offset; + pendingResult + .getInputChannelOffsets() + .computeIfAbsent( + seg.channelInfo(), unused -> new StateContentMetaInfo()) + .withDataAdded(offset, size); + NetworkActionsLogger.tracePersist( + action, + seg.length() + " bytes", + seg.channelInfo(), + checkpointId); + } + }); + } finally { + try { + reader.close(); + } catch (Exception e) { + LOG.info( + "Failed to close the fetched channel state reader of checkpoint {}", + checkpointId, + e); + } + } + } + void writeOutput( JobVertexID jobVertexID, int subtaskIndex, ResultSubpartitionInfo info, Buffer buffer) { try { diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateSerializer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateSerializer.java index 252d25c2e29..ec858460dd8 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateSerializer.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateSerializer.java @@ -43,6 +43,8 @@ interface ChannelStateSerializer { void writeData(DataOutputStream stream, Buffer... flinkBuffers) throws IOException; + void writeData(DataOutputStream stream, InputStream input, int length) throws IOException; + void readHeader(InputStream stream) throws IOException; int readLength(InputStream stream) throws IOException; @@ -165,6 +167,18 @@ class ChannelStateSerializerImpl implements ChannelStateSerializer { } } + @Override + public void writeData(DataOutputStream stream, InputStream input, int length) + throws IOException { + Preconditions.checkArgument(length >= 0, "negative state size"); + stream.writeInt(length); + long copied = input.transferTo(stream); + if (copied != length) { + throw new java.io.EOFException( + "Unexpected EOF: expected " + length + " bytes of segment body, got " + copied); + } + } + private int getSize(Buffer[] buffers) { int len = 0; for (Buffer buffer : buffers) { diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequest.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequest.java index abef241c325..d1913df0416 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequest.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequest.java @@ -258,6 +258,20 @@ abstract class ChannelStateWriteRequest { return new CheckpointAbortRequest(jobVertexID, subtaskIndex, checkpointId, cause); } + static ChannelStateWriteRequest replayInputDataFromSpill( + JobVertexID jobVertexID, + int subtaskIndex, + long checkpointId, + FetchedChannelStateReader reader) { + return new CheckpointInProgressRequest( + "writeInputFromSpill", + jobVertexID, + subtaskIndex, + checkpointId, + writer -> writer.writeInputFromSpill(jobVertexID, subtaskIndex, reader), + throwable -> reader.close()); + } + static ChannelStateWriteRequest registerSubtask(JobVertexID jobVertexID, int subtaskIndex) { return new SubtaskRegisterRequest(jobVertexID, subtaskIndex); } diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriter.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriter.java index 6fee1402036..d10d88ea800 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriter.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriter.java @@ -25,6 +25,9 @@ import org.apache.flink.runtime.state.InputChannelStateHandle; import org.apache.flink.runtime.state.ResultSubpartitionStateHandle; import org.apache.flink.util.CloseableIterator; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + import java.io.Closeable; import java.util.Collection; import java.util.Collections; @@ -190,10 +193,15 @@ public interface ChannelStateWriter extends Closeable { ChannelStateWriteResult getAndRemoveWriteResult(long checkpointId) throws IllegalArgumentException; + /** Records input-channel state from a spill file and takes ownership of {@code reader}. */ + void addInputDataFromSpill(long checkpointId, FetchedChannelStateReader reader); + ChannelStateWriter NO_OP = new NoOpChannelStateWriter(); /** No-op implementation of {@link ChannelStateWriter}. */ class NoOpChannelStateWriter implements ChannelStateWriter { + private static final Logger LOG = LoggerFactory.getLogger(NoOpChannelStateWriter.class); + @Override public void start(long checkpointId, CheckpointOptions checkpointOptions) {} @@ -231,6 +239,18 @@ public interface ChannelStateWriter extends Closeable { CompletableFuture.completedFuture(Collections.emptyList())); } + @Override + public void addInputDataFromSpill(long checkpointId, FetchedChannelStateReader reader) { + try { + reader.close(); + } catch (Exception e) { + LOG.info( + "Failed to close the fetched channel state reader of checkpoint {}", + checkpointId, + e); + } + } + @Override public void close() {} } diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriterImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriterImpl.java index 40d7ddffd1e..21db97355db 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriterImpl.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriterImpl.java @@ -42,6 +42,7 @@ import java.util.concurrent.atomic.AtomicBoolean; import static org.apache.flink.runtime.checkpoint.channel.ChannelStateWriteRequest.completeInput; import static org.apache.flink.runtime.checkpoint.channel.ChannelStateWriteRequest.completeOutput; +import static org.apache.flink.runtime.checkpoint.channel.ChannelStateWriteRequest.replayInputDataFromSpill; import static org.apache.flink.runtime.checkpoint.channel.ChannelStateWriteRequest.write; /** @@ -235,6 +236,12 @@ public class ChannelStateWriterImpl implements ChannelStateWriter { enqueue(completeOutput(jobVertexID, subtaskIndex, checkpointId), false); } + @Override + public void addInputDataFromSpill(long checkpointId, FetchedChannelStateReader reader) { + LOG.debug("{} replaying input data from spill, checkpoint {}", taskName, checkpointId); + enqueue(replayInputDataFromSpill(jobVertexID, subtaskIndex, checkpointId, reader), false); + } + @Override public void abort(long checkpointId, Throwable cause, boolean cleanup) { LOG.debug("{} aborting, checkpoint {}", taskName, checkpointId); diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/checkpointing/ChannelState.java b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/checkpointing/ChannelState.java index b0dfb4291a1..31288059fd3 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/checkpointing/ChannelState.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/checkpointing/ChannelState.java @@ -19,6 +19,7 @@ package org.apache.flink.streaming.runtime.io.checkpointing; import org.apache.flink.runtime.checkpoint.CheckpointException; +import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter; import org.apache.flink.runtime.checkpoint.channel.FetchedChannelStateSnapshot; import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo; import org.apache.flink.runtime.checkpoint.channel.RecoveryCheckpointTrigger; @@ -54,14 +55,19 @@ final class ChannelState { private final RecoveryCheckpointTrigger recoveryCheckpointTrigger; + private final ChannelStateWriter channelStateWriter; + public ChannelState(CheckpointableInput[] inputs) { - this(inputs, RecoveryCheckpointTrigger.NO_OP); + this(inputs, RecoveryCheckpointTrigger.NO_OP, ChannelStateWriter.NO_OP); } public ChannelState( - CheckpointableInput[] inputs, RecoveryCheckpointTrigger recoveryCheckpointTrigger) { + CheckpointableInput[] inputs, + RecoveryCheckpointTrigger recoveryCheckpointTrigger, + ChannelStateWriter channelStateWriter) { this.inputs = inputs; this.recoveryCheckpointTrigger = checkNotNull(recoveryCheckpointTrigger); + this.channelStateWriter = checkNotNull(channelStateWriter); } public void blockChannel(InputChannelInfo channelInfo) { @@ -113,19 +119,19 @@ final class ChannelState { } /** - * Dispatches checkpoint start: inserts recovery-checkpoint barriers into in-recovery channels - * through the trigger, then notifies every input. (FLINK-38544 transitional: the spilling - * backend adds a third step handing a reader opened from the snapshot to the channel-state - * writer, instead of closing it here.) + * Transfers spill-snapshot ownership to the writer after all inputs observe checkpoint start. */ public void onCheckpointStartedForAllInputs(CheckpointBarrier barrier) throws CheckpointException, IOException { long cpId = barrier.getId(); + // The snapshot is closed either way: closing is a no-op once the reader below was opened + // (the writer owns and closes it), and releases the grant if we never got that far. try (FetchedChannelStateSnapshot snapshot = recoveryCheckpointTrigger.snapshotAndInsertBarriers(cpId)) { for (CheckpointableInput input : inputs) { input.checkpointStarted(barrier); } + channelStateWriter.addInputDataFromSpill(cpId, snapshot.reader()); } } } diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/checkpointing/InputProcessorUtil.java b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/checkpointing/InputProcessorUtil.java index 39f85fe08b2..e299399e993 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/checkpointing/InputProcessorUtil.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/checkpointing/InputProcessorUtil.java @@ -22,6 +22,7 @@ import org.apache.flink.api.common.operators.MailboxExecutor; import org.apache.flink.configuration.CheckpointingOptions; import org.apache.flink.configuration.Configuration; import org.apache.flink.core.execution.CheckpointingMode; +import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter; import org.apache.flink.runtime.checkpoint.channel.RecoveryCheckpointTrigger; import org.apache.flink.runtime.io.network.partition.consumer.CheckpointableInput; import org.apache.flink.runtime.io.network.partition.consumer.IndexedInputGate; @@ -123,6 +124,7 @@ public class InputProcessorUtil { Clock clock = SystemClock.getInstance(); CheckpointingMode checkpointingMode = CheckpointingOptions.getCheckpointingMode(jobConf); + ChannelStateWriter channelStateWriter = checkpointCoordinator.getChannelStateWriter(); switch (checkpointingMode) { case EXACTLY_ONCE: int numberOfChannels = @@ -141,7 +143,8 @@ public class InputProcessorUtil { inputs, clock, numberOfChannels, - recoveryCheckpointTrigger); + recoveryCheckpointTrigger, + channelStateWriter); case AT_LEAST_ONCE: if (CheckpointingOptions.isUnalignedCheckpointEnabled(jobConf)) { throw new IllegalStateException( @@ -175,7 +178,8 @@ public class InputProcessorUtil { CheckpointableInput[] inputs, Clock clock, int numberOfChannels, - RecoveryCheckpointTrigger recoveryCheckpointTrigger) { + RecoveryCheckpointTrigger recoveryCheckpointTrigger, + ChannelStateWriter channelStateWriter) { boolean enableCheckpointAfterTasksFinished = config.getConfiguration() .get(CheckpointingOptions.ENABLE_CHECKPOINTS_AFTER_TASKS_FINISH); @@ -189,6 +193,7 @@ public class InputProcessorUtil { BarrierAlignmentUtil.createRegisterTimerCallback(mailboxExecutor, timerService), enableCheckpointAfterTasksFinished, recoveryCheckpointTrigger, + channelStateWriter, inputs); } else { return SingleCheckpointBarrierHandler.aligned( diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/checkpointing/SingleCheckpointBarrierHandler.java b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/checkpointing/SingleCheckpointBarrierHandler.java index 8594fba540f..547941c2447 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/checkpointing/SingleCheckpointBarrierHandler.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/io/checkpointing/SingleCheckpointBarrierHandler.java @@ -22,6 +22,7 @@ import org.apache.flink.annotation.Internal; import org.apache.flink.annotation.VisibleForTesting; import org.apache.flink.runtime.checkpoint.CheckpointException; import org.apache.flink.runtime.checkpoint.CheckpointFailureReason; +import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter; import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo; import org.apache.flink.runtime.checkpoint.channel.RecoveryCheckpointTrigger; import org.apache.flink.runtime.io.network.api.CancelCheckpointMarker; @@ -121,6 +122,7 @@ public class SingleCheckpointBarrierHandler extends CheckpointBarrierHandler { }, enableCheckpointsAfterTasksFinish, RecoveryCheckpointTrigger.NO_OP, + ChannelStateWriter.NO_OP, inputs); } @@ -133,6 +135,7 @@ public class SingleCheckpointBarrierHandler extends CheckpointBarrierHandler { DelayableTimer registerTimer, boolean enableCheckpointAfterTasksFinished, RecoveryCheckpointTrigger recoveryCheckpointTrigger, + ChannelStateWriter channelStateWriter, CheckpointableInput... inputs) { return new SingleCheckpointBarrierHandler( taskName, @@ -141,7 +144,8 @@ public class SingleCheckpointBarrierHandler extends CheckpointBarrierHandler { clock, numOpenChannels, new AlternatingWaitingForFirstBarrierUnaligned( - false, new ChannelState(inputs, recoveryCheckpointTrigger)), + false, + new ChannelState(inputs, recoveryCheckpointTrigger, channelStateWriter)), false, registerTimer, inputs, @@ -178,6 +182,7 @@ public class SingleCheckpointBarrierHandler extends CheckpointBarrierHandler { DelayableTimer registerTimer, boolean enableCheckpointAfterTasksFinished, RecoveryCheckpointTrigger recoveryCheckpointTrigger, + ChannelStateWriter channelStateWriter, CheckpointableInput... inputs) { return new SingleCheckpointBarrierHandler( taskName, @@ -186,7 +191,7 @@ public class SingleCheckpointBarrierHandler extends CheckpointBarrierHandler { clock, numOpenChannels, new AlternatingWaitingForFirstBarrier( - new ChannelState(inputs, recoveryCheckpointTrigger)), + new ChannelState(inputs, recoveryCheckpointTrigger, channelStateWriter)), true, registerTimer, inputs, diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriterImplTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriterImplTest.java index fb931946bbb..50ec1cc6afb 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriterImplTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriterImplTest.java @@ -30,11 +30,14 @@ import org.apache.flink.runtime.state.storage.JobManagerCheckpointStorage; import org.apache.flink.util.function.BiConsumerWithException; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; import java.io.IOException; import java.util.ArrayDeque; +import java.util.Collections; import java.util.Deque; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; import static org.apache.flink.util.CloseableIterator.ofElements; @@ -51,6 +54,8 @@ class ChannelStateWriterImplTest { private static final CheckpointStorage CHECKPOINT_STORAGE = new JobManagerCheckpointStorage(); + @TempDir private java.nio.file.Path tempDir; + @Test void testAddEventBuffer() throws Exception { @@ -371,10 +376,124 @@ class ChannelStateWriterImplTest { writer.finishInput(CHECKPOINT_ID); writer.finishOutput(CHECKPOINT_ID); } + + @Test + void testAddInputDataFromSpillAsyncDemux() throws Exception { + SyncChannelStateWriteRequestExecutor worker = + new SyncChannelStateWriteRequestExecutor(JOB_ID); + try (ChannelStateWriterImpl writer = newWriter(worker)) { + worker.registerSubtask(JOB_VERTEX_ID, SUBTASK_INDEX); + writer.start(CHECKPOINT_ID, CheckpointOptions.forCheckpointWithDefaultLocation()); + + InputChannelInfo c0 = new InputChannelInfo(0, 0); + InputChannelInfo c1 = new InputChannelInfo(0, 1); + + FetchedChannelState state; + try (TestSpillWriter spillWriter = new TestSpillWriter(tempDir)) { + spillWriter.writeRecord(c0, new byte[] {1, 2, 3}, 3); + spillWriter.writeRecord(c1, new byte[] {4, 5}, 2); + spillWriter.writeRecord(c0, new byte[] {6}, 1); + state = spillWriter.getChannelState(); + } + FetchedChannelStateReader reader = state.reader(); + // Drop the handoff grant; the reader now holds the only outstanding grant. + state.release(); + + writer.addInputDataFromSpill(CHECKPOINT_ID, reader); + // Request is queued but not yet processed — state must still be alive. + assertThat(state.isClosed()).isFalse(); + + worker.processAllRequests(); + // After processing, the reader is closed by the request, releasing the last grant. + assertThat(state.isClosed()).isTrue(); + } + } + + @Test + void testAddInputDataFromSpillEmptySnapshotStillSubmitted() throws Exception { + // Empty readers are no longer short-circuited; they are submitted to the writer thread. + QueueCountingExecutor worker = new QueueCountingExecutor(); + try (ChannelStateWriterImpl writer = + new ChannelStateWriterImpl( + JOB_VERTEX_ID, + TASK_NAME, + SUBTASK_INDEX, + new ConcurrentHashMap<>(), + worker, + 5)) { + worker.registerSubtask(JOB_VERTEX_ID, SUBTASK_INDEX); + writer.start(CHECKPOINT_ID, CheckpointOptions.forCheckpointWithDefaultLocation()); + + int submittedBefore = worker.submitCount.get(); + FetchedChannelState emptyState = new FetchedChannelState(Collections.emptyList()); + FetchedChannelStateReader emptyReader = emptyState.reader(); + emptyState.release(); + + writer.addInputDataFromSpill(CHECKPOINT_ID, emptyReader); + + assertThat(worker.submitCount.get()) + .as("empty reader must still be submitted to the writer thread") + .isGreaterThan(submittedBefore); + } + } + + @Test + void testAddInputDataFromSpillSegmentsClosedOnSuccess() throws Exception { + SyncChannelStateWriteRequestExecutor worker = + new SyncChannelStateWriteRequestExecutor(JOB_ID); + try (ChannelStateWriterImpl writer = newWriter(worker)) { + worker.registerSubtask(JOB_VERTEX_ID, SUBTASK_INDEX); + writer.start(CHECKPOINT_ID, CheckpointOptions.forCheckpointWithDefaultLocation()); + + FetchedChannelState state; + try (TestSpillWriter spillWriter = new TestSpillWriter(tempDir)) { + spillWriter.writeRecord(new InputChannelInfo(0, 0), new byte[] {1}, 1); + state = spillWriter.getChannelState(); + } + FetchedChannelStateReader reader = state.reader(); + state.release(); + + writer.addInputDataFromSpill(CHECKPOINT_ID, reader); + worker.processAllRequests(); + + // After processing, the last grant is released and the state is cleaned up. + assertThat(state.isClosed()).isTrue(); + } + } + + private ChannelStateWriterImpl newWriter(SyncChannelStateWriteRequestExecutor worker) { + return new ChannelStateWriterImpl( + JOB_VERTEX_ID, TASK_NAME, SUBTASK_INDEX, new ConcurrentHashMap<>(), worker, 5); + } } class TestException extends RuntimeException {} +/** Counts submissions without processing them, for the empty-snapshot submission test. */ +class QueueCountingExecutor implements ChannelStateWriteRequestExecutor { + + final AtomicInteger submitCount = new AtomicInteger(0); + + @Override + public void submit(ChannelStateWriteRequest e) { + submitCount.incrementAndGet(); + } + + @Override + public void submitPriority(ChannelStateWriteRequest e) { + submitCount.incrementAndGet(); + } + + @Override + public void start() throws IllegalStateException {} + + @Override + public void registerSubtask(JobVertexID jobVertexID, int subtaskIndex) {} + + @Override + public void releaseSubtask(JobVertexID jobVertexID, int subtaskIndex) {} +} + class SyncChannelStateWriteRequestExecutor implements ChannelStateWriteRequestExecutor { private final ChannelStateWriteRequestDispatcher requestProcessor; private final Deque<ChannelStateWriteRequest> deque; diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/MockChannelStateWriter.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/MockChannelStateWriter.java index c77208f3ff7..bdff6d44718 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/MockChannelStateWriter.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/MockChannelStateWriter.java @@ -74,6 +74,16 @@ public class MockChannelStateWriter implements ChannelStateWriter { } } + @Override + public void addInputDataFromSpill(long checkpointId, FetchedChannelStateReader reader) { + checkCheckpointId(checkpointId); + try { + reader.close(); + } catch (Exception e) { + rethrow(e); + } + } + @Override public void addOutputData( long checkpointId, ResultSubpartitionInfo info, int startSeqNum, Buffer... data) { diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/io/checkpointing/ChannelStateTest.java b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/io/checkpointing/ChannelStateTest.java index f6dc7f7f69b..0307a92f1af 100644 --- a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/io/checkpointing/ChannelStateTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/io/checkpointing/ChannelStateTest.java @@ -21,29 +21,34 @@ package org.apache.flink.streaming.runtime.io.checkpointing; import org.apache.flink.runtime.checkpoint.CheckpointException; import org.apache.flink.runtime.checkpoint.CheckpointOptions; import org.apache.flink.runtime.checkpoint.CheckpointType; +import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter; import org.apache.flink.runtime.checkpoint.channel.FetchedChannelState; +import org.apache.flink.runtime.checkpoint.channel.FetchedChannelStateReader; import org.apache.flink.runtime.checkpoint.channel.FetchedChannelStateSnapshot; import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo; import org.apache.flink.runtime.checkpoint.channel.RecoveryCheckpointTrigger; +import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo; import org.apache.flink.runtime.io.network.api.CheckpointBarrier; +import org.apache.flink.runtime.io.network.buffer.Buffer; import org.apache.flink.runtime.io.network.partition.consumer.CheckpointableInput; import org.apache.flink.runtime.state.CheckpointStorageLocationReference; +import org.apache.flink.util.CloseableIterator; import org.junit.jupiter.api.Test; import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; import static org.assertj.core.api.Assertions.assertThat; /** * Verifies the {@link ChannelState#onCheckpointStartedForAllInputs} dispatcher: call ordering and * feature-off no-op routing through the {@link RecoveryCheckpointTrigger#NO_OP} singleton. - * - * <p>FLINK-38544 transitional: this covers the 2-step dispatch (trigger, then per-input - * notification); the spilling backend adds a third step handing the trigger's snapshot reader to - * the channel-state writer and completes this test to cover all three steps. */ class ChannelStateTest { @@ -52,11 +57,15 @@ class ChannelStateTest { @Test void testStepOrderingFeatureOn() throws Exception { List<String> trace = new ArrayList<>(); - RecordingTrigger trigger = new RecordingTrigger(trace); + // An empty reader is sufficient to verify ordering. + FetchedChannelStateSnapshot snap = FetchedChannelState.emptySnapshot(); + RecordingTrigger trigger = new RecordingTrigger(trace, snap); + RecordingWriter writer = new RecordingWriter(trace); CheckpointableInput input1 = new RecordingInput(trace, "in1"); CheckpointableInput input2 = new RecordingInput(trace, "in2"); - ChannelState state = new ChannelState(new CheckpointableInput[] {input1, input2}, trigger); + ChannelState state = + new ChannelState(new CheckpointableInput[] {input1, input2}, trigger, writer); CheckpointBarrier barrier = newUnalignedBarrier(); state.onCheckpointStartedForAllInputs(barrier); @@ -65,21 +74,49 @@ class ChannelStateTest { .containsExactly( "trigger.snapshotAndInsertBarriers:" + CHECKPOINT_ID, "in1.checkpointStarted:" + CHECKPOINT_ID, - "in2.checkpointStarted:" + CHECKPOINT_ID); + "in2.checkpointStarted:" + CHECKPOINT_ID, + "writer.addInputDataFromSpill:" + CHECKPOINT_ID); } @Test void testStepOrderingFeatureOff() throws Exception { List<String> trace = new ArrayList<>(); + RecordingWriter writer = new RecordingWriter(trace); CheckpointableInput input = new RecordingInput(trace, "in1"); ChannelState state = new ChannelState( - new CheckpointableInput[] {input}, RecoveryCheckpointTrigger.NO_OP); + new CheckpointableInput[] {input}, RecoveryCheckpointTrigger.NO_OP, writer); + + state.onCheckpointStartedForAllInputs(newUnalignedBarrier()); + + assertThat(trace) + .containsExactly( + "in1.checkpointStarted:" + CHECKPOINT_ID, + "writer.addInputDataFromSpill:" + CHECKPOINT_ID); + assertThat(writer.lastSnapshotWasEmpty.get()).isTrue(); + } + + @Test + void testEmptySnapshotStillSubmitted() throws Exception { + // Empty readers (no spill files) are no longer short-circuited; they still reach + // addInputDataFromSpill on the writer thread. + List<String> trace = new ArrayList<>(); + FetchedChannelStateSnapshot emptySnap = FetchedChannelState.emptySnapshot(); + RecordingTrigger trigger = new RecordingTrigger(trace, emptySnap); + RecordingWriter writer = new RecordingWriter(trace); + + ChannelState state = + new ChannelState( + new CheckpointableInput[] {new RecordingInput(trace, "in1")}, + trigger, + writer); state.onCheckpointStartedForAllInputs(newUnalignedBarrier()); - assertThat(trace).containsExactly("in1.checkpointStarted:" + CHECKPOINT_ID); + // Empty reader must still reach the writer (no inline short-circuit). + assertThat(writer.addInputDataFromSpillCalls.get()).isEqualTo(1); + assertThat(writer.lastSnapshotWasEmpty.get()).isTrue(); } private static CheckpointBarrier newUnalignedBarrier() { @@ -93,18 +130,83 @@ class ChannelStateTest { private static final class RecordingTrigger implements RecoveryCheckpointTrigger { private final List<String> trace; + private final FetchedChannelStateSnapshot snapshot; - RecordingTrigger(List<String> trace) { + RecordingTrigger(List<String> trace, FetchedChannelStateSnapshot snapshot) { this.trace = trace; + this.snapshot = snapshot; } @Override public FetchedChannelStateSnapshot snapshotAndInsertBarriers(long checkpointId) { trace.add("trigger.snapshotAndInsertBarriers:" + checkpointId); - return FetchedChannelState.emptySnapshot(); + return snapshot; } } + private static final class RecordingWriter implements ChannelStateWriter { + private final List<String> trace; + final AtomicBoolean lastSnapshotWasEmpty = new AtomicBoolean(false); + final AtomicLong lastCpId = new AtomicLong(-1L); + final AtomicInteger addInputDataFromSpillCalls = new AtomicInteger(0); + + RecordingWriter(List<String> trace) { + this.trace = trace; + } + + @Override + public void start(long checkpointId, CheckpointOptions checkpointOptions) {} + + @Override + public void addInputData( + long checkpointId, + InputChannelInfo info, + int startSeqNum, + CloseableIterator<Buffer> data) {} + + @Override + public void addOutputData( + long checkpointId, ResultSubpartitionInfo info, int startSeqNum, Buffer... data) {} + + @Override + public void addOutputDataFuture( + long checkpointId, + ResultSubpartitionInfo info, + int startSeqNum, + CompletableFuture<List<Buffer>> data) {} + + @Override + public void finishInput(long checkpointId) {} + + @Override + public void finishOutput(long checkpointId) {} + + @Override + public void abort(long checkpointId, Throwable cause, boolean cleanup) {} + + @Override + public ChannelStateWriteResult getAndRemoveWriteResult(long checkpointId) { + return ChannelStateWriteResult.EMPTY; + } + + @Override + public void addInputDataFromSpill(long checkpointId, FetchedChannelStateReader reader) { + trace.add("writer.addInputDataFromSpill:" + checkpointId); + lastCpId.set(checkpointId); + addInputDataFromSpillCalls.incrementAndGet(); + try { + // Peek whether the reader has any segments by attempting the first advance. + // The first nextSegment() call is exempt from the "previous body consumed" rule. + lastSnapshotWasEmpty.set(reader.advanceAndGetNextSegment().isEmpty()); + reader.close(); + } catch (Exception ignored) { + } + } + + @Override + public void close() {} + } + private static final class RecordingInput implements CheckpointableInput { private final List<String> trace; diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/io/checkpointing/TestBarrierHandlerFactory.java b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/io/checkpointing/TestBarrierHandlerFactory.java index dca71e35f6b..8c281c04a6b 100644 --- a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/io/checkpointing/TestBarrierHandlerFactory.java +++ b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/io/checkpointing/TestBarrierHandlerFactory.java @@ -74,6 +74,7 @@ public class TestBarrierHandlerFactory { actionRegistration, enableCheckpointsAfterTasksFinish, RecoveryCheckpointTrigger.NO_OP, + stateWriter, inputGate); } }
