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 0b5cb41c4298540a3d206447c25c6d59b5fe69ad Author: Rui Fan <[email protected]> AuthorDate: Mon Jul 6 01:59:11 2026 +0200 [FLINK-39524][checkpoint] Forward-only spill reader with snapshot/resume FetchedChannelStateReader interface: nextSegment(), snapshot(), emptyReader(); inner SpillSegment (channelInfo(), bodyStream(), length(), commit()). FetchedChannelStateReaderImpl: sequential file IO over the spill files, a bounded per-segment body InputStream, current/committed positions, and snapshot resume with skip-only-on-first-positioning. Completes the pieces deferred from the container commit for dependency reasons: FetchedChannelStateSnapshot (the immutable one-shot resume point holding one lifecycle grant and a Position; reader() at most once, fail-loud), FetchedChannelState#reader(), and FetchedChannelStateRefCountTest#testReaderAcquiresAndReleasesOnClose. All files in this commit are byte-identical to their final form. --- .../checkpoint/channel/FetchedChannelState.java | 25 +- .../channel/FetchedChannelStateDrainer.java | 6 +- .../channel/FetchedChannelStateReader.java | 107 ++++ .../channel/FetchedChannelStateReaderImpl.java | 444 ++++++++++++++ .../channel/FetchedChannelStateSnapshot.java | 156 +++++ .../channel/RecoveryCheckpointTrigger.java | 17 +- .../runtime/io/checkpointing/ChannelState.java | 13 +- .../flink/streaming/runtime/tasks/StreamTask.java | 2 +- .../channel/FetchedChannelStateReaderTest.java | 637 +++++++++++++++++++++ .../runtime/io/checkpointing/ChannelStateTest.java | 5 +- 10 files changed, 1389 insertions(+), 23 deletions(-) diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelState.java index 1137578a2cf..3cef621f29b 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelState.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelState.java @@ -40,15 +40,12 @@ import static org.apache.flink.util.Preconditions.checkNotNull; * table is maintained. The reader scans files sequentially, reading each 12-byte header to obtain * the channel info and body length. * - * <p>The file list grows as the writer rotates to new files (one rotation per 64 MB soft limit), - * and is sealed on writer close. + * <p>The file list is fixed at construction: rotation to new files (one per 64 MB soft limit) + * happens earlier, in the writer, which builds this container from the final list on close. * * <p>File lifecycle is managed by {@link #acquire()} / {@link #release()} reference counting. Files * are deleted only when the last lifecycle grant is released (i.e. when both the main reader and * all snapshot readers have finished). - * - * <p>Mutations (file list appends) are single-writer and intentionally unsynchronized; callers must - * serialize them via the channel IO executor. */ @Internal public final class FetchedChannelState implements Closeable { @@ -74,6 +71,24 @@ public final class FetchedChannelState implements Closeable { // Read-phase API (called by the reader after the writer is sealed) // ------------------------------------------------------------------------------------------- + /** + * Opens the main reader covering all segments from the beginning. The returned reader holds one + * lifecycle grant and must be closed when done. + */ + public FetchedChannelStateReader reader() { + return snapshotAtStart().reader(); + } + + /** A snapshot covering all segments from the beginning. */ + public FetchedChannelStateSnapshot snapshotAtStart() { + return new FetchedChannelStateSnapshot(this, 0, 0L, null, 0); + } + + /** A snapshot over no spill files: its reader yields no segments. */ + public static FetchedChannelStateSnapshot emptySnapshot() { + return new FetchedChannelState(Collections.emptyList()).snapshotAtStart(); + } + /** Returns the ordered list of spill file paths. Read-only view. */ public List<Path> files() { return Collections.unmodifiableList(files); diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateDrainer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateDrainer.java index 9d5a84d8330..23f33bca40e 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateDrainer.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateDrainer.java @@ -71,10 +71,14 @@ public final class FetchedChannelStateDrainer implements RecoveryCheckpointTrigg * must persist is already inside the channels' queues, so the snapshot is inherently empty. */ @Override - public void snapshotAndInsertBarriers(long checkpointId) throws IOException { + public FetchedChannelStateSnapshot snapshotAndInsertBarriers(long checkpointId) + throws IOException { for (RecoverableInputChannel channel : channels) { channel.insertRecoveryCheckpointBarrierIfInRecovery(checkpointId); } + // No snapshot side for the in-memory backend: everything a checkpoint must persist is + // already inside the channels' queues, so the returned snapshot is empty. + return FetchedChannelState.emptySnapshot(); } @Override diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateReader.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateReader.java new file mode 100644 index 00000000000..0e507c2eb0d --- /dev/null +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateReader.java @@ -0,0 +1,107 @@ +/* + * 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.channel; + +import org.apache.flink.annotation.Internal; + +import java.io.Closeable; +import java.io.InputStream; +import java.util.Optional; + +/** + * Forward-only, strictly sequential reader over a {@link FetchedChannelState}'s spill files. + * Deliberately not a Java {@link java.util.Iterator}: a body must be fully read before advancing, + * body ownership is handed to the consumer, and consume/commit are separate steps. + * + * <p>The main reader (opened via {@link FetchedChannelState#reader()}, starting at offset 0) + * records the delivered boundary — the "committed position" — via {@link SpillSegment#commit()}; + * before anything is committed it equals the reader's start position. Each checkpoint derives a + * {@link #snapshot()} that resumes from that boundary. {@link #snapshot()} and {@link + * SpillSegment#commit()} must be called under the drainer lock; disk reads happen outside it. + */ +@Internal +public interface FetchedChannelStateReader extends Closeable { + + /** + * Advances to the next segment and returns it, or {@link Optional#empty()} when no segment + * remains. + * + * <p>Entry rule (the first call is exempt): the previous segment's body must be fully read, + * otherwise this is a contract violation and fails loud (no skip-ahead). + */ + Optional<SpillSegment> advanceAndGetNextSegment(); + + /** + * Derives an independent resume point starting from the committed position. The snapshot holds + * its own {@link FetchedChannelState} lifecycle grant; the caller must open a reader from it + * via {@link FetchedChannelStateSnapshot#reader()} and close that reader when done. + * + * <p>Must be called under the drainer lock so that the copied position reflects the latest + * committed state. + * + * @return a snapshot capturing the current committed position; caller must open and close a + * reader from it + */ + FetchedChannelStateSnapshot snapshot(); + + /** + * One per-channel segment produced by {@link #advanceAndGetNextSegment()}. + * + * <p>The segment body bytes are opaque to the reader; record framing is handled by the + * consumer's deserializer. A consumer reads {@link #bodyStream()} to EOF (after {@link + * #length()} bytes), and the drain consumer additionally calls {@link #commit()}. + * + * <p>Ownership of {@link #bodyStream()} passes to the consumer: the reader no longer tracks how + * far it has been read. The "previous body must be fully read" rule (no skip-ahead) is enforced + * at the next {@link FetchedChannelStateReader#advanceAndGetNextSegment()} call, not here. + * + * <p>A segment is valid only until the next {@code advanceAndGetNextSegment()} call on the + * parent reader. + */ + interface SpillSegment { + + /** The channel whose data this segment contains. */ + InputChannelInfo channelInfo(); + + /** + * Returns an {@link InputStream} bounded to this segment's body. Reading returns {@code -1} + * (EOF) after {@link #length()} bytes; it never reads into the next segment or the next + * file. + * + * <p>The stream is single-use, not thread-safe, and must be fully consumed before the next + * {@link FetchedChannelStateReader#advanceAndGetNextSegment()}. + */ + InputStream bodyStream(); + + /** + * Number of body bytes this segment hands out before EOF. For the snapshot path this is the + * not-yet-delivered remainder used as the length prefix when writing to the checkpoint + * stream. Bounded by the spill file size limit, so it always fits in an {@code int}. + */ + int length(); + + /** + * Records the body bytes read from {@link #bodyStream()} so far as delivered. + * + * <p>Called once per delivered buffer by the drainer's drain loop, under the same lock as + * {@link FetchedChannelStateReader#snapshot()}, so a snapshot always resumes on a buffer + * boundary. Only the main reader commits. + */ + void commit(); + } +} diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateReaderImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateReaderImpl.java new file mode 100644 index 00000000000..e9dc1782679 --- /dev/null +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateReaderImpl.java @@ -0,0 +1,444 @@ +/* + * 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.channel; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.runtime.checkpoint.channel.FetchedChannelStateReader.SpillSegment; +import org.apache.flink.util.IOUtils; + +import javax.annotation.Nullable; + +import java.io.BufferedInputStream; +import java.io.ByteArrayInputStream; +import java.io.DataInputStream; +import java.io.EOFException; +import java.io.IOException; +import java.io.InputStream; +import java.nio.channels.Channels; +import java.nio.channels.SeekableByteChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.List; +import java.util.Optional; + +import static org.apache.flink.runtime.checkpoint.channel.AbstractSpillingHandler.SEGMENT_HEADER_BYTES; +import static org.apache.flink.util.Preconditions.checkState; + +/** + * The single {@link FetchedChannelStateReader} implementation over a {@link FetchedChannelState}'s + * spill files. + * + * <p>Reading is strictly sequential: the stream is positioned once when a file is opened and then + * only moves forward. A snapshot that resumes inside a segment body starts right at the boundary — + * the segment's channel and remaining length come from the snapshot itself, so no header is re-read + * and no prefix is skipped. + * + * <p>The reader tracks two {@link Position}s of the same shape — the values a {@link + * FetchedChannelStateSnapshot} is made of: {@code current}, the live read position, and {@code + * committed}, the delivered boundary that {@link SpillSegment#commit()} publishes from it under the + * drainer lock. + * + * <p>The "previous body fully read before advancing" rule is checked at the {@link + * #advanceAndGetNextSegment()} entry. Body ownership is handed to the consumer, so the reader does + * not track body progress except through {@code current}. + */ +@Internal +final class FetchedChannelStateReaderImpl implements FetchedChannelStateReader { + + private final FetchedChannelStateSnapshot snapshot; + + /** + * The live read position: the file being read and the exact byte offset within it. It runs + * ahead of the committed boundary by what the drainer has read from disk but not yet handed to + * a channel (at most one recovery buffer) — reads happen outside the drainer lock, delivery and + * commit inside it. A checkpoint therefore cuts at the committed boundary: what is already in + * the channel queue is persisted from the channel side, everything after it is re-read from the + * spill files. + */ + private final Position current; + + private final Position committed; + + /** Open stream over {@code current.fileIndex}, or {@code null} before the first read. */ + @Nullable private InputStream currentFileStream; + + /** Size of the file currently open. */ + private long currentFileSize; + + /** + * The segment handed out by the last {@link #advanceAndGetNextSegment()}. A segment is only + * valid until the next one is handed out; reading or committing an older one fails loud. + */ + @Nullable private Segment currentSegment; + + private boolean positioned; + private boolean closed; + + FetchedChannelStateReaderImpl(FetchedChannelStateSnapshot snapshot) { + this.snapshot = snapshot; + this.current = + new Position( + snapshot.fileIndex(), + snapshot.readOffset(), + snapshot.channel(), + snapshot.remaining()); + this.committed = current.copy(); + } + + @Override + public Optional<SpillSegment> advanceAndGetNextSegment() { + checkState(!closed, "FetchedChannelStateReader is closed"); + checkState( + !positioned || current.remaining == 0, + "Previous segment body not fully consumed before advancing: %s bytes left", + current.remaining); + try { + if (!positioned) { + positioned = true; + if (current.channel != null) { + return resumedSegment(); + } + } + return nextSegmentAtHeader(); + } catch (IOException e) { + throw new RuntimeException("Failed to read segment", e); + } + } + + /** + * Resume path, taken once by a snapshot reader whose start offset sits inside a body: the + * channel and the remaining length come from the snapshot, so nothing is re-read or skipped. + */ + private Optional<SpillSegment> resumedSegment() throws IOException { + openFileAndSeek(); + currentSegment = + new Segment( + this, + current.channel, + new BoundedSegmentStream(currentFileStream, current, current.remaining)); + return Optional.of(currentSegment); + } + + /** Steady path: the stream sits on a segment header; read it and hand out the whole body. */ + private Optional<SpillSegment> nextSegmentAtHeader() throws IOException { + if (!openCurrentFile()) { + return Optional.empty(); + } + SegmentHeader header = readHeaderAtCurrent(); + current.startSegment(header.channelInfo, header.bufferLength); + if (currentSegment != null) { + currentSegment.body.invalidate(); + } + currentSegment = + new Segment( + this, + header.channelInfo, + new BoundedSegmentStream(currentFileStream, current, header.bufferLength)); + return Optional.of(currentSegment); + } + + @Override + public FetchedChannelStateSnapshot snapshot() { + checkState(!closed, "FetchedChannelStateReader is closed"); + return new FetchedChannelStateSnapshot( + snapshot.channelState(), + committed.fileIndex, + committed.readOffset, + committed.channel, + committed.remaining); + } + + @Override + public void close() throws IOException { + if (closed) { + return; + } + closed = true; + try { + closeFileStream(); + } finally { + snapshot.release(); + } + } + + // ------------------------------------------------------------------------------------------- + // Sequential IO over the spill files; all of it advances current.readOffset / current.fileIndex + // ------------------------------------------------------------------------------------------- + + /** The spill files in write order; the reader never mutates the list. */ + private List<Path> files() { + return snapshot.channelState().files(); + } + + /** + * Ensures a file is open with the stream positioned at {@code current}'s read offset, ready to + * read this segment's header. Rolls to the next file when the current one is exhausted. Returns + * false when no segment remains. + */ + private boolean openCurrentFile() throws IOException { + boolean rolled = false; + while (current.fileIndex < files().size()) { + openFileAndSeek(); + if (current.readOffset < currentFileSize) { + return true; + } + // Current file fully read: move to the next file's first segment. The writer never + // produces an empty file, so one roll always lands on data. + checkState(!rolled, "Rolled past more than one empty file"); + closeFileStream(); + current.rollToNextFile(); + rolled = true; + } + return false; + } + + /** Reads the 12-byte header at the current read offset; advances past it. */ + private SegmentHeader readHeaderAtCurrent() throws IOException { + byte[] headerBytes = new byte[SEGMENT_HEADER_BYTES]; + readFully(headerBytes); + DataInputStream h = new DataInputStream(new ByteArrayInputStream(headerBytes)); + int gateIdx = h.readInt(); + int channelIdx = h.readInt(); + int bufferLength = h.readInt(); + checkState(bufferLength >= 0, "negative segment length: %s", bufferLength); + checkState( + gateIdx >= 0 && channelIdx >= 0, + "negative channel info in segment header: %s/%s", + gateIdx, + channelIdx); + return new SegmentHeader(new InputChannelInfo(gateIdx, channelIdx), bufferLength); + } + + /** + * Ensures the file at {@code current.fileIndex} is open with the stream positioned at {@code + * current.readOffset}. If a stream is already open it is left as-is: sequential reading + * guarantees it is already there. + */ + private void openFileAndSeek() throws IOException { + if (currentFileStream != null) { + return; + } + SeekableByteChannel channel = + Files.newByteChannel(files().get(current.fileIndex), StandardOpenOption.READ); + try { + currentFileSize = channel.size(); + channel.position(current.readOffset); + } catch (IOException e) { + channel.close(); + throw e; + } + currentFileStream = new BufferedInputStream(Channels.newInputStream(channel)); + } + + private void readFully(byte[] buf) throws IOException { + IOUtils.readFully(currentFileStream, buf, 0, buf.length); + current.advanceReadOffset(buf.length); + } + + private void closeFileStream() throws IOException { + if (currentFileStream != null) { + currentFileStream.close(); + currentFileStream = null; + } + } + + // ------------------------------------------------------------------------------------------- + // Position: where the open stream sits + // ------------------------------------------------------------------------------------------- + + /** + * A point in the spill files, in the four values a {@link FetchedChannelStateSnapshot} is made + * of: {@code channel} is null exactly when {@code readOffset} sits on a segment header, and + * non-null while a body is in flight, {@code remaining} being what is left of it. + */ + static final class Position { + private int fileIndex; + private long readOffset; + @Nullable private InputChannelInfo channel; + private int remaining; + + Position( + int fileIndex, long readOffset, @Nullable InputChannelInfo channel, int remaining) { + this.fileIndex = fileIndex; + this.readOffset = readOffset; + this.channel = channel; + this.remaining = remaining; + } + + Position copy() { + return new Position(fileIndex, readOffset, channel, remaining); + } + + /** Publishes {@code other} into this position (used by commit). */ + void copyFrom(Position other) { + fileIndex = other.fileIndex; + readOffset = other.readOffset; + channel = other.channel; + remaining = other.remaining; + } + + /** Advances past the {@code delta} header bytes just read; no body is in flight. */ + void advanceReadOffset(long delta) { + readOffset += delta; + } + + /** Enters the body of a freshly read segment header. */ + void startSegment(InputChannelInfo segmentChannel, int bodyLength) { + channel = segmentChannel; + remaining = bodyLength; + } + + /** Accounts for {@code n} body bytes handed to the consumer. */ + void advanceBody(int n) { + readOffset += n; + remaining -= n; + if (remaining == 0) { + channel = null; + } + } + + /** Rolls to the start of the next file once the current one is exhausted. */ + void rollToNextFile() { + fileIndex++; + readOffset = 0L; + } + } + + /** Parsed segment header: channel and full body length. */ + private static final class SegmentHeader { + private final InputChannelInfo channelInfo; + private final int bufferLength; + + private SegmentHeader(InputChannelInfo channelInfo, int bufferLength) { + this.channelInfo = channelInfo; + this.bufferLength = bufferLength; + } + } + + /** + * The single {@link SpillSegment} implementation. Exposes one segment's channel, body, and + * length; {@link #commit()} advances the reader's {@code committed} position to however many + * body bytes have been read. Reading the body and committing are separate steps so the consumer + * can read outside the drainer lock and commit inside it. + * + * <p>Only the main reader commits. + */ + private static final class Segment implements SpillSegment { + private final FetchedChannelStateReaderImpl reader; + private final InputChannelInfo channelInfo; + private final BoundedSegmentStream body; + + private Segment( + FetchedChannelStateReaderImpl reader, + InputChannelInfo channelInfo, + BoundedSegmentStream body) { + this.reader = reader; + this.channelInfo = channelInfo; + this.body = body; + } + + @Override + public InputChannelInfo channelInfo() { + return channelInfo; + } + + @Override + public InputStream bodyStream() { + return body; + } + + @Override + public int length() { + return body.deliverableLength(); + } + + @Override + public void commit() { + checkState( + reader.currentSegment == this, + "Committing a segment that is no longer the current one"); + reader.committed.copyFrom(reader.current); + } + } + + /** + * A forward-only, bounded view over the body bytes this reader still has to hand out for one + * segment. The bound is {@code current.remaining}, so the view keeps no counter of its own; it + * reaches EOF there and never reads into the next segment or file. If the file ends first, an + * {@link EOFException} is thrown (fail-loud). Closing this view does not close the underlying + * file; the reader owns it. + */ + private static final class BoundedSegmentStream extends InputStream { + private final InputStream fileStream; + private final Position current; + private final int length; + + /** Set when the reader hands out the next segment; this view must not be read after. */ + private boolean stale; + + private BoundedSegmentStream(InputStream fileStream, Position current, int length) { + this.fileStream = fileStream; + this.current = current; + this.length = length; + } + + private void invalidate() { + stale = true; + } + + /** Number of body bytes this view will hand out. */ + int deliverableLength() { + return length; + } + + @Override + public int read() throws IOException { + byte[] one = new byte[1]; + int n = read(one, 0, 1); + return n < 0 ? -1 : (one[0] & 0xFF); + } + + @Override + public int read(byte[] buf, int off, int len) throws IOException { + checkState(!stale, "Reading a segment that is no longer the current one"); + if (current.remaining == 0) { + return -1; + } + int toRead = Math.min(len, current.remaining); + int n = fileStream.read(buf, off, toRead); + if (n > 0) { + current.advanceBody(n); + } + if (n < 0) { + throw new EOFException( + "Unexpected EOF in segment body after " + + (length - current.remaining) + + "/" + + length + + " bytes"); + } + return n; + } + + @Override + public void close() { + // Do not close the underlying file; it is owned by the reader. + } + } +} diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateSnapshot.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateSnapshot.java new file mode 100644 index 00000000000..8585d0c8d55 --- /dev/null +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateSnapshot.java @@ -0,0 +1,156 @@ +/* + * 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.channel; + +import org.apache.flink.annotation.Internal; + +import javax.annotation.Nullable; + +import java.io.IOException; + +import static org.apache.flink.util.Preconditions.checkArgument; +import static org.apache.flink.util.Preconditions.checkState; + +/** + * An immutable resume point for a {@link FetchedChannelState} reader: where to start reading, plus + * the metadata of a half-delivered segment when the boundary falls inside one. It holds one + * lifecycle grant on the underlying {@link FetchedChannelState} (acquired in the constructor). + * + * <p>{@code channel} decides how the resume point is read. It is {@code null} when {@code + * readOffset} points at a segment header, and non-null when it points into a segment body, in which + * case {@code remaining} is that segment's not-yet-delivered byte count. The three cases that occur + * in practice: + * + * <pre> + * (0, 0, null, 0) nothing delivered yet — read everything from the first header + * (2, 840, null, 0) boundary landed exactly on a header — read that header, then continue + * (2, 851, c3, 57) boundary landed inside c3's body — hand out its last 57 bytes, no header read + * </pre> + * + * <p>A snapshot is a one-shot, single-reader handle: exactly one {@link FetchedChannelStateReader} + * may be opened from it via {@link #reader()}, and opening one transfers the grant to that reader, + * which returns it on close. A second {@link #reader()} call fails loud. + * + * <p>Owners must {@link #close()} the snapshot. Closing releases the grant when no reader was + * opened; once one was, the reader owns it and closing here is a no-op — so closing early never + * deletes files out from under a live reader. + */ +@Internal +public final class FetchedChannelStateSnapshot implements AutoCloseable { + + private final FetchedChannelState channelState; + + /** Spill file to resume in. */ + private final int fileIndex; + + /** Byte offset within that file to resume at. */ + private final long readOffset; + + /** Channel of the half-delivered segment, or {@code null} if {@code readOffset} is a header. */ + @Nullable private final InputChannelInfo channel; + + /** Not-yet-delivered body bytes of that segment; 0 iff {@code channel} is {@code null}. */ + private final int remaining; + + /** True once {@link #reader()} has been called; prevents opening a second reader. */ + private boolean readerOpened; + + private boolean closed; + + /** + * Creates a snapshot resuming at {@code readOffset} in file {@code fileIndex}. Acquires one + * lifecycle grant on {@code channelState}; the grant is released when the reader returned by + * {@link #reader()} is closed. + */ + FetchedChannelStateSnapshot( + FetchedChannelState channelState, + int fileIndex, + long readOffset, + @Nullable InputChannelInfo channel, + int remaining) { + checkArgument( + (channel == null) == (remaining == 0), + "channel and remaining must be set together: %s / %s", + channel, + remaining); + this.channelState = channelState; + this.fileIndex = fileIndex; + this.readOffset = readOffset; + this.channel = channel; + this.remaining = remaining; + channelState.acquire(); + } + + /** + * Opens the reader for this snapshot. May be called at most once; a second call fails loud to + * enforce the 1:1 snapshot-to-reader invariant. + * + * @return a new reader starting from this snapshot's position; caller must close it when done + */ + public FetchedChannelStateReader reader() { + checkState(!closed, "Snapshot is closed"); + checkState(!readerOpened, "A reader has already been opened from this snapshot"); + readerOpened = true; + return new FetchedChannelStateReaderImpl(this); + } + + /** + * Releases the lifecycle grant held by this snapshot. Called by the reader on close; must not + * be called directly by any other party. + */ + void release() throws IOException { + channelState.release(); + } + + /** + * Releases the grant if no reader was opened; otherwise the reader owns it and this is a no-op. + * Idempotent. + */ + @Override + public void close() throws IOException { + if (closed) { + return; + } + closed = true; + if (!readerOpened) { + channelState.release(); + } + } + + /** Returns the underlying channel state (package-private; used by the reader). */ + FetchedChannelState channelState() { + return channelState; + } + + int fileIndex() { + return fileIndex; + } + + long readOffset() { + return readOffset; + } + + @Nullable + InputChannelInfo channel() { + return channel; + } + + int remaining() { + return remaining; + } +} diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveryCheckpointTrigger.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveryCheckpointTrigger.java index 1ef35d4f407..58f724abe8b 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveryCheckpointTrigger.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveryCheckpointTrigger.java @@ -27,18 +27,15 @@ import java.io.IOException; public interface RecoveryCheckpointTrigger { /** - * Atomically snapshots the undrained recovered state and inserts matching {@link - * RecoveryCheckpointBarrier}s into in-recovery channels. - * - * <p>FLINK-38544 transitional signature: once the spilling backend lands this returns an - * independent reader over the remaining (undrained) spill segments that the caller owns and - * must close. The in-memory backend hands all recovered buffers to the channels in one shot, so - * there is never an undrained residue and nothing to return yet. + * Atomically snapshots the undrained spill slice and inserts matching {@link + * RecoveryCheckpointBarrier}s into in-recovery channels. Returns a snapshot over the remaining + * segments; the caller owns it, opens the reader it needs and must close the snapshot. */ - void snapshotAndInsertBarriers(long checkpointId) throws IOException, CheckpointException; + FetchedChannelStateSnapshot snapshotAndInsertBarriers(long checkpointId) + throws IOException, CheckpointException; - /** Inserts no barriers (and there is no recovered state left to snapshot). */ - RecoveryCheckpointTrigger NO_OP = checkpointId -> {}; + /** Returns an empty snapshot (no spill files, so no segments) and inserts no barriers. */ + RecoveryCheckpointTrigger NO_OP = checkpointId -> FetchedChannelState.emptySnapshot(); RecoveryCheckpointTrigger NOT_READY = ign -> { 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 503e1532724..b0dfb4291a1 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.FetchedChannelStateSnapshot; import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo; import org.apache.flink.runtime.checkpoint.channel.RecoveryCheckpointTrigger; import org.apache.flink.runtime.io.network.api.CheckpointBarrier; @@ -114,15 +115,17 @@ 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 the trigger's snapshot reader to the channel-state writer.) + * backend adds a third step handing a reader opened from the snapshot to the channel-state + * writer, instead of closing it here.) */ public void onCheckpointStartedForAllInputs(CheckpointBarrier barrier) throws CheckpointException, IOException { long cpId = barrier.getId(); - recoveryCheckpointTrigger.snapshotAndInsertBarriers(cpId); - - for (CheckpointableInput input : inputs) { - input.checkpointStarted(barrier); + try (FetchedChannelStateSnapshot snapshot = + recoveryCheckpointTrigger.snapshotAndInsertBarriers(cpId)) { + for (CheckpointableInput input : inputs) { + input.checkpointStarted(barrier); + } } } } diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java index 8c9270f581b..22cd3050ce2 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java @@ -1119,7 +1119,7 @@ public abstract class StreamTask<OUT, OP extends StreamOperator<OUT>> public RecoveryCheckpointTrigger getRecoveryCheckpointTrigger() { return cpId -> { checkState(mailboxProcessor.isMailboxThread()); - recoveryCheckpointTrigger.snapshotAndInsertBarriers(cpId); + return recoveryCheckpointTrigger.snapshotAndInsertBarriers(cpId); }; } diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateReaderTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateReaderTest.java new file mode 100644 index 00000000000..2f424078c2d --- /dev/null +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateReaderTest.java @@ -0,0 +1,637 @@ +/* + * 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.channel; + +import org.apache.flink.runtime.checkpoint.channel.FetchedChannelStateReader.SpillSegment; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.EOFException; +import java.io.IOException; +import java.io.InputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Optional; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Unit tests for {@link FetchedChannelStateReader}: sequential segment scanning, body boundedness, + * cross-file transparency, snapshot derivation, and fail-loud on truncated segments. + * + * <p>Segment boundaries are self-described in disk headers; no in-memory segment locator table is + * used. + */ +class FetchedChannelStateReaderTest { + + @TempDir Path tempDir; + + // ------------------------------------------------------------------------------------------- + // Segment iteration + // ------------------------------------------------------------------------------------------- + + @Test + void testIteratorEmptyWhenNoDataWritten() throws Exception { + // A writer that never spills produces no state; an empty state has no segments. + FetchedChannelState state = new FetchedChannelState(Collections.emptyList()); + try (FetchedChannelStateReader reader = state.reader()) { + assertThat(reader.advanceAndGetNextSegment()).isEmpty(); + } + } + + @Test + void testMultipleIteratorIteratedInOrder() throws Exception { + InputChannelInfo c0 = new InputChannelInfo(0, 0); + InputChannelInfo c1 = new InputChannelInfo(0, 1); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writeRecord(c0, bytes(10), 1); + writer.writeRecord(c1, bytes(20), 1); + writer.writeRecord(c0, bytes(30), 1); + state = writer.getChannelState(); + } + + List<InputChannelInfo> channels = new ArrayList<>(); + try (FetchedChannelStateReader reader = state.reader()) { + Optional<SpillSegment> next; + while ((next = reader.advanceAndGetNextSegment()).isPresent()) { + SpillSegment seg = next.get(); + channels.add(seg.channelInfo()); + readAll(seg.bodyStream()); + } + } + + // Segments are produced at channel switches: c0, c1, c0 + assertThat(channels).containsExactly(c0, c1, c0); + } + + // ------------------------------------------------------------------------------------------- + // Body boundedness: body() stops exactly at segment end + // ------------------------------------------------------------------------------------------- + + @Test + void testBodyReturnsMinus1AtSegmentEnd() throws Exception { + InputChannelInfo ch = new InputChannelInfo(0, 0); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writeRecord(ch, bytes(1, 2), 2); + state = writer.getChannelState(); + } + + try (FetchedChannelStateReader reader = state.reader()) { + SpillSegment seg = reader.advanceAndGetNextSegment().orElseThrow(AssertionError::new); + InputStream body = seg.bodyStream(); + // Read exactly length bytes + byte[] data = new byte[seg.length()]; + int totalRead = 0; + while (totalRead < data.length) { + int n = body.read(data, totalRead, data.length - totalRead); + assertThat(n).isGreaterThan(0); + totalRead += n; + } + // Next read must return EOF + assertThat(body.read()).isEqualTo(-1); + } + } + + // ------------------------------------------------------------------------------------------- + // Cross-file transparency + // ------------------------------------------------------------------------------------------- + + @Test + void testCrossFileTransparencyWhenRotationOccurs() throws Exception { + InputChannelInfo c0 = new InputChannelInfo(0, 0); + InputChannelInfo c1 = new InputChannelInfo(0, 1); + + // Use tiny rotation threshold so first segment triggers a file rotation. + FetchedChannelState state; + try (TestSpillWriter writer = + new TestSpillWriter( + tempDir, + 1L /* 1 byte file bound */, + AbstractSpillingHandler.DEFAULT_MAX_SEGMENT_SIZE_BYTES)) { + writer.writeRecord(c0, bytes(10, 11, 12), 3); + writer.writeRecord(c1, bytes(20, 21), 2); + state = writer.getChannelState(); + } + + // Two segments in different files. + assertThat(state.files()).hasSize(2); + + List<InputChannelInfo> channels = new ArrayList<>(); + try (FetchedChannelStateReader reader = state.reader()) { + Optional<SpillSegment> next; + while ((next = reader.advanceAndGetNextSegment()).isPresent()) { + SpillSegment seg = next.get(); + channels.add(seg.channelInfo()); + // Body read must not throw even if the segment is in a different file. + readAll(seg.bodyStream()); + } + } + + assertThat(channels).containsExactly(c0, c1); + } + + // ------------------------------------------------------------------------------------------- + // Snapshot: independent reader with correct start position + // ------------------------------------------------------------------------------------------- + + @Test + void testSnapshotCoversAllIteratorWhenNothingConsumed() throws Exception { + InputChannelInfo c0 = new InputChannelInfo(0, 0); + InputChannelInfo c1 = new InputChannelInfo(0, 1); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writeRecord(c0, bytes(1), 1); + writer.writeRecord(c1, bytes(2), 1); + state = writer.getChannelState(); + } + + try (FetchedChannelStateReader root = state.reader()) { + // Snapshot before consuming anything + try (FetchedChannelStateReader snap = root.snapshot().reader()) { + List<InputChannelInfo> channels = new ArrayList<>(); + Optional<SpillSegment> next; + while ((next = snap.advanceAndGetNextSegment()).isPresent()) { + SpillSegment seg = next.get(); + channels.add(seg.channelInfo()); + readAll(seg.bodyStream()); + } + assertThat(channels).containsExactly(c0, c1); + } + } + } + + @Test + void testSnapshotAfterFullSegmentConsumedSkipsThatSegment() throws Exception { + InputChannelInfo c0 = new InputChannelInfo(0, 0); + InputChannelInfo c1 = new InputChannelInfo(0, 1); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writeRecord(c0, bytes(1), 1); + writer.writeRecord(c1, bytes(2), 1); + state = writer.getChannelState(); + } + + try (FetchedChannelStateReader root = state.reader()) { + // Consume and commit first segment + SpillSegment first = root.advanceAndGetNextSegment().orElseThrow(AssertionError::new); + readAll(first.bodyStream()); + first.commit(); + + // Snapshot must start from second segment + try (FetchedChannelStateReader snap = root.snapshot().reader()) { + List<InputChannelInfo> channels = new ArrayList<>(); + Optional<SpillSegment> next; + while ((next = snap.advanceAndGetNextSegment()).isPresent()) { + SpillSegment seg = next.get(); + channels.add(seg.channelInfo()); + readAll(seg.bodyStream()); + } + assertThat(channels).containsExactly(c1); + } + } + } + + @Test + void testSnapshotFromMidSegmentStartsAtCommittedByteOffset() throws Exception { + InputChannelInfo ch = new InputChannelInfo(0, 0); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + // Two records in the same channel -> one segment + writer.writeRecord(ch, bytes(10, 11), 2); + state = writer.getChannelState(); + } + // Verify: one file with one segment + assertThat(state.files()).hasSize(1); + + try (FetchedChannelStateReader root = state.reader()) { + SpillSegment seg = root.advanceAndGetNextSegment().orElseThrow(AssertionError::new); + int fullLength = seg.length(); + InputStream body = seg.bodyStream(); + + // Read only 1 byte without committing, then snapshot — snapshot should start from 0 + // (no bytes committed yet). + body.read(); + + try (FetchedChannelStateReader snapBeforeCommit = root.snapshot().reader()) { + SpillSegment snapSeg = + snapBeforeCommit + .advanceAndGetNextSegment() + .orElseThrow(AssertionError::new); + assertThat(snapSeg.length()).isEqualTo(fullLength); + readAll(snapSeg.bodyStream()); + } + + // Read rest of body and commit + readAll(body); + seg.commit(); + + // After commit the snapshot must be empty + try (FetchedChannelStateReader snapAfterCommit = root.snapshot().reader()) { + assertThat(snapAfterCommit.advanceAndGetNextSegment()).isEmpty(); + } + } + } + + @Test + void testSnapshotAfterPartialCommitReadsRemainingBodyTail() throws Exception { + InputChannelInfo ch = new InputChannelInfo(0, 0); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + // One segment whose body is a single pass-through blob of known bytes. + writer.writePassThrough(ch, bytes(1, 2, 3, 4, 5, 6, 7, 8), 0, 8); + state = writer.getChannelState(); + } + assertThat(state.files()).hasSize(1); + + try (FetchedChannelStateReader root = state.reader()) { + SpillSegment seg = root.advanceAndGetNextSegment().orElseThrow(AssertionError::new); + int fullLength = seg.length(); + InputStream body = seg.bodyStream(); + + // Read and commit only a 3-byte prefix. + byte[] prefix = new byte[3]; + assertThat(body.read(prefix)).isEqualTo(3); + seg.commit(); + + // Snapshot must resume mid-segment and yield exactly the remaining tail bytes. + try (FetchedChannelStateReader snap = root.snapshot().reader()) { + SpillSegment snapSeg = + snap.advanceAndGetNextSegment().orElseThrow(AssertionError::new); + assertThat(snapSeg.channelInfo()).isEqualTo(ch); + assertThat(snapSeg.length()).isEqualTo(fullLength - 3); + byte[] tail = readAll(snapSeg.bodyStream()); + assertThat(tail).isEqualTo(bytes(4, 5, 6, 7, 8)); + assertThat(snap.advanceAndGetNextSegment()).isEmpty(); + } + } + } + + @Test + void testSnapshotCarriesResumeSegmentOnlyWhenBoundaryIsMidBody() throws Exception { + InputChannelInfo ch = new InputChannelInfo(0, 0); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writePassThrough(ch, bytes(1, 2, 3, 4, 5, 6, 7, 8), 0, 8); + state = writer.getChannelState(); + } + + try (FetchedChannelStateReader root = state.reader()) { + // Before anything is committed: start of the first file, no resume segment. + try (FetchedChannelStateSnapshot atStart = root.snapshot()) { + assertThat(atStart.fileIndex()).isZero(); + assertThat(atStart.readOffset()).isZero(); + assertThat(atStart.channel()).isNull(); + assertThat(atStart.remaining()).isZero(); + } + + SpillSegment seg = root.advanceAndGetNextSegment().orElseThrow(AssertionError::new); + InputStream body = seg.bodyStream(); + assertThat(body.read(new byte[3])).isEqualTo(3); + seg.commit(); + + // Boundary inside the body: the channel and its remainder travel with the snapshot. + try (FetchedChannelStateSnapshot midBody = root.snapshot()) { + assertThat(midBody.channel()).isEqualTo(ch); + assertThat(midBody.remaining()).isEqualTo(5); + } + + readAll(body); + seg.commit(); + + // Boundary on a header (here: end of the last segment): no resume segment again. + try (FetchedChannelStateSnapshot atHeader = root.snapshot()) { + assertThat(atHeader.channel()).isNull(); + assertThat(atHeader.remaining()).isZero(); + } + } + } + + @Test + void testSnapshotResumesPartialSegmentAcrossFileBoundary() throws Exception { + InputChannelInfo c0 = new InputChannelInfo(0, 0); + InputChannelInfo c1 = new InputChannelInfo(0, 1); + + // Tiny rotation threshold so the two segments land in separate files. + FetchedChannelState state; + try (TestSpillWriter writer = + new TestSpillWriter( + tempDir, 1L, AbstractSpillingHandler.DEFAULT_MAX_SEGMENT_SIZE_BYTES)) { + writer.writePassThrough(c0, bytes(1, 2, 3, 4), 0, 4); + writer.writePassThrough(c1, bytes(5, 6, 7), 0, 3); + state = writer.getChannelState(); + } + assertThat(state.files()).hasSize(2); + + try (FetchedChannelStateReader root = state.reader()) { + SpillSegment first = root.advanceAndGetNextSegment().orElseThrow(AssertionError::new); + // Commit a 1-byte prefix of the file-0 segment. + first.bodyStream().read(new byte[1]); + first.commit(); + + try (FetchedChannelStateReader snap = root.snapshot().reader()) { + SpillSegment resumed = + snap.advanceAndGetNextSegment().orElseThrow(AssertionError::new); + assertThat(resumed.channelInfo()).isEqualTo(c0); + assertThat(readAll(resumed.bodyStream())).isEqualTo(bytes(2, 3, 4)); + + // Crossing into file 1 must reset the skip to 0. + SpillSegment following = + snap.advanceAndGetNextSegment().orElseThrow(AssertionError::new); + assertThat(following.channelInfo()).isEqualTo(c1); + assertThat(readAll(following.bodyStream())).isEqualTo(bytes(5, 6, 7)); + + assertThat(snap.advanceAndGetNextSegment()).isEmpty(); + } + } + } + + @Test + void testRootDrainViaRepeatedCommitsTerminatesAndFinalSnapshotEmpty() throws Exception { + InputChannelInfo c0 = new InputChannelInfo(0, 0); + InputChannelInfo c1 = new InputChannelInfo(0, 1); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writePassThrough(c0, bytes(1, 2), 0, 2); + writer.writePassThrough(c1, bytes(3, 4, 5), 0, 3); + state = writer.getChannelState(); + } + + try (FetchedChannelStateReader root = state.reader()) { + int count = 0; + Optional<SpillSegment> next; + while ((next = root.advanceAndGetNextSegment()).isPresent()) { + SpillSegment seg = next.get(); + readAll(seg.bodyStream()); + seg.commit(); + count++; + } + assertThat(count).isEqualTo(2); + + // After draining everything, a snapshot must have nothing left. + try (FetchedChannelStateReader snap = root.snapshot().reader()) { + assertThat(snap.advanceAndGetNextSegment()).isEmpty(); + } + } + } + + // ------------------------------------------------------------------------------------------- + // Fail-loud on truncated segment + // ------------------------------------------------------------------------------------------- + + @Test + void testBodyThrowsEOFExceptionOnTruncatedFile() throws Exception { + InputChannelInfo ch = new InputChannelInfo(0, 0); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writeRecord(ch, bytes(1, 2, 3, 4, 5, 6, 7, 8), 8); + state = writer.getChannelState(); + } + + // Truncate the spill file to just the header (12 bytes) so the body is missing. + Path spill = state.files().get(0); + byte[] headerOnly = Files.readAllBytes(spill); + // Keep only the 12-byte header, discard body. + Files.write( + spill, + java.util.Arrays.copyOf(headerOnly, AbstractSpillingHandler.SEGMENT_HEADER_BYTES), + StandardOpenOption.TRUNCATE_EXISTING); + + try (FetchedChannelStateReader reader = state.reader()) { + SpillSegment seg = reader.advanceAndGetNextSegment().orElseThrow(AssertionError::new); + // bufferLength from header says > 0 bytes, but file has nothing after the header. + assertThatThrownBy(() -> readAll(seg.bodyStream())).isInstanceOf(EOFException.class); + } + } + + // ------------------------------------------------------------------------------------------- + // Reference counting: acquire/release via reader lifecycle + // ------------------------------------------------------------------------------------------- + + @Test + void testReaderAcquiresAndReleasesRefCount() throws Exception { + InputChannelInfo ch = new InputChannelInfo(0, 0); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writeRecord(ch, bytes(1), 1); + state = writer.getChannelState(); + } + + Path spill = state.files().get(0); + + FetchedChannelStateReader reader = state.reader(); + // Drop the handoff grant so the reader's grant is the only one outstanding. + state.release(); + assertThat(java.nio.file.Files.exists(spill)).isTrue(); + + reader.close(); + + // After closing the only reader, the file is cleaned up. + assertThat(java.nio.file.Files.exists(spill)).isFalse(); + } + + @Test + void testSnapshotClosedWithoutReaderReleasesItsGrant() throws Exception { + InputChannelInfo ch = new InputChannelInfo(0, 0); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writeRecord(ch, bytes(1), 1); + state = writer.getChannelState(); + } + + Path spill = state.files().get(0); + + FetchedChannelStateReader root = state.reader(); + FetchedChannelStateSnapshot snapshot = root.snapshot(); + // Drop the handoff grant so only the reader's and the snapshot's grants remain. + state.release(); + + root.close(); // The snapshot still holds a grant, so the file survives. + assertThat(Files.exists(spill)).isTrue(); + + snapshot.close(); // No reader was ever opened: closing returns the snapshot's own grant. + assertThat(Files.exists(spill)).isFalse(); + + snapshot.close(); // Idempotent: no second release, no exception. + assertThatThrownBy(snapshot::reader).isInstanceOf(IllegalStateException.class); + } + + @Test + void testClosingSnapshotAfterOpeningReaderDoesNotDeleteFilesEarly() throws Exception { + InputChannelInfo ch = new InputChannelInfo(0, 0); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writeRecord(ch, bytes(1), 1); + state = writer.getChannelState(); + } + + Path spill = state.files().get(0); + + FetchedChannelStateReader root = state.reader(); + FetchedChannelStateSnapshot snapshot = root.snapshot(); + FetchedChannelStateReader snapReader = snapshot.reader(); + state.release(); + root.close(); + + // The grant moved to the reader, so closing the snapshot must not drop it. + snapshot.close(); + assertThat(Files.exists(spill)).isTrue(); + + snapReader.close(); + assertThat(Files.exists(spill)).isFalse(); + } + + @Test + void testSnapshotKeepsFilesAliveUntilSnapshotClosed() throws Exception { + InputChannelInfo ch = new InputChannelInfo(0, 0); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writeRecord(ch, bytes(1), 1); + state = writer.getChannelState(); + } + + Path spill = state.files().get(0); + + FetchedChannelStateReader root = state.reader(); + FetchedChannelStateReader snap = root.snapshot().reader(); + // Drop the handoff grant so only the two reader grants remain outstanding. + state.release(); + + root.close(); // One grant released; file must still exist because snap holds another. + assertThat(java.nio.file.Files.exists(spill)).isTrue(); + + snap.close(); // Last grant released; file must be deleted. + assertThat(java.nio.file.Files.exists(spill)).isFalse(); + } + + // ------------------------------------------------------------------------------------------- + // New behaviour: first-call exemption, fail-loud on body not consumed, empty reader + // ------------------------------------------------------------------------------------------- + + @Test + void testFirstNextSegmentCallDoesNotRequirePreviousBodyConsumed() throws Exception { + // The "previous body must be fully consumed" rule must not fire on the very first call + // because there is no previous segment. + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writeRecord(new InputChannelInfo(0, 0), bytes(1, 2), 2); + state = writer.getChannelState(); + } + + try (FetchedChannelStateReader reader = state.reader()) { + // Must not throw on the first call even though no body has been consumed before. + Optional<SpillSegment> seg = reader.advanceAndGetNextSegment(); + assertThat(seg).isPresent(); + } + } + + @Test + void testNextSegmentThrowsWhenPreviousBodyNotFullyConsumed() throws Exception { + InputChannelInfo ch = new InputChannelInfo(0, 0); + InputChannelInfo ch2 = new InputChannelInfo(0, 1); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writeRecord(ch, bytes(1, 2, 3, 4), 4); + writer.writeRecord(ch2, bytes(5, 6), 2); + state = writer.getChannelState(); + } + + try (FetchedChannelStateReader reader = state.reader()) { + SpillSegment seg = reader.advanceAndGetNextSegment().orElseThrow(AssertionError::new); + // Read only part of the body — do not exhaust it. + seg.bodyStream().read(); + + // Advancing to the next segment while the previous body is not fully consumed must fail + // loud. + assertThatThrownBy(reader::advanceAndGetNextSegment) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("Previous segment body not fully consumed"); + } + } + + @Test + void testEmptyReaderNextSegmentReturnsEmptyAndCloseIsClean() throws Exception { + FetchedChannelState state = new FetchedChannelState(Collections.emptyList()); + try (FetchedChannelStateReader reader = state.reader()) { + // First call on an empty reader must return empty without throwing. + assertThat(reader.advanceAndGetNextSegment()).isEmpty(); + // Closing must not throw. + } + } + + @Test + void testEmptySnapshotHandsOutIndependentReadersSoCloseDoesNotLeak() throws Exception { + // An empty snapshot is created and closed once per checkpoint. close() is single-use (it + // flips the closed flag permanently), so each call must yield a fresh instance; otherwise + // the first consumer's close would make every later consumer's advanceAndGetNextSegment() + // fail loud. + FetchedChannelStateReader first = FetchedChannelState.emptySnapshot().reader(); + assertThat(first.advanceAndGetNextSegment()).isEmpty(); + first.close(); + + FetchedChannelStateReader second = FetchedChannelState.emptySnapshot().reader(); + assertThat(second).isNotSameAs(first); + // Must still work after the previously obtained empty reader was closed. + assertThat(second.advanceAndGetNextSegment()).isEmpty(); + second.close(); + } + + // ------------------------------------------------------------------------------------------- + // Helpers + // ------------------------------------------------------------------------------------------- + + private static byte[] readAll(InputStream in) throws IOException { + java.io.ByteArrayOutputStream out = new java.io.ByteArrayOutputStream(); + byte[] buf = new byte[256]; + int n; + while ((n = in.read(buf)) != -1) { + out.write(buf, 0, n); + } + return out.toByteArray(); + } + + private static byte[] bytes(int... values) { + byte[] arr = new byte[values.length]; + for (int i = 0; i < values.length; i++) { + arr[i] = (byte) values[i]; + } + return arr; + } +} 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 b0e21aef390..f6dc7f7f69b 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,6 +21,8 @@ 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.FetchedChannelState; +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.io.network.api.CheckpointBarrier; @@ -97,8 +99,9 @@ class ChannelStateTest { } @Override - public void snapshotAndInsertBarriers(long checkpointId) { + public FetchedChannelStateSnapshot snapshotAndInsertBarriers(long checkpointId) { trace.add("trigger.snapshotAndInsertBarriers:" + checkpointId); + return FetchedChannelState.emptySnapshot(); } }
