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 227805ba7bf625643a45da44a480939ad49ef0b7 Author: Rui Fan <[email protected]> AuthorDate: Mon Jul 6 01:52:04 2026 +0200 [FLINK-39524][checkpoint] Spill-writing handlers: segmented on-disk format New AbstractSpillingHandler (in RecoveredChannelStateHandler.java, extending the abstract input-channel base): a reusable DataOutputSerializer accumulates one channel's segment [gateIdx][channelIdx][bufferLength][body...]; the body length is backfilled via writeIntUnsafe at seal; flush through OffsetAwareOutputStream; 64 MiB soft rotation; closeInternal() seals, closes the stream, and builds the FetchedChannelState handoff (calling acquire()). New SpillingNoFilteringHandler writes the recovered buffer's bytes verbatim via segmentSerializerFor(...).write(memorySegment, offset, len). flink-core OffsetAwareOutputStream: ctor package-private -> public so flink-runtime can construct it. Also adds the @Nullable FetchedChannelState getProducedChannelState() hook to AbstractInputChannelRecoveredStateHandler (returns null; the spilling base overrides it) and restores the final class javadoc on the abstract base. The factory is NOT switched: no production path selects a spilling handler yet. Transitional deviation: AbstractSpillingHandler passes the base's transitional third ctor arg (true); reverts to the 2-arg base ctor when the spilling backend lands. FetchedChannelStateRefCountTest rides here (it needs TestSpillWriter) minus testReaderAcquiresAndReleasesOnClose, which follows with the reader commit. --- .../flink/core/fs/OffsetAwareOutputStream.java | 2 +- .../channel/RecoveredChannelStateHandler.java | 281 ++++++++++++++++++++- .../channel/AbstractSpillingHandlerTest.java | 184 ++++++++++++++ .../checkpoint/channel/TestSpillWriter.java | 105 ++++++++ 4 files changed, 570 insertions(+), 2 deletions(-) diff --git a/flink-core/src/main/java/org/apache/flink/core/fs/OffsetAwareOutputStream.java b/flink-core/src/main/java/org/apache/flink/core/fs/OffsetAwareOutputStream.java index 3ee4b761e1b..375c95da25a 100644 --- a/flink-core/src/main/java/org/apache/flink/core/fs/OffsetAwareOutputStream.java +++ b/flink-core/src/main/java/org/apache/flink/core/fs/OffsetAwareOutputStream.java @@ -35,7 +35,7 @@ public final class OffsetAwareOutputStream implements Closeable { private long position; - OffsetAwareOutputStream(OutputStream currentOut, long position) { + public OffsetAwareOutputStream(OutputStream currentOut, long position) { this.currentOut = checkNotNull(currentOut); this.position = position; } diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandler.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandler.java index 6a68f47c639..98660963afe 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandler.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandler.java @@ -18,6 +18,8 @@ package org.apache.flink.runtime.checkpoint.channel; import org.apache.flink.annotation.VisibleForTesting; +import org.apache.flink.core.fs.OffsetAwareOutputStream; +import org.apache.flink.core.memory.DataOutputSerializer; import org.apache.flink.core.memory.MemorySegment; import org.apache.flink.core.memory.MemorySegmentFactory; import org.apache.flink.runtime.checkpoint.InflightDataRescalingDescriptor; @@ -38,13 +40,21 @@ import org.apache.flink.runtime.io.network.partition.consumer.RecoveredInputChan import javax.annotation.Nonnull; import javax.annotation.Nullable; +import java.io.BufferedOutputStream; import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.UUID; import static org.apache.flink.runtime.checkpoint.channel.ChannelStateByteBuffer.wrap; import static org.apache.flink.util.Preconditions.checkArgument; +import static org.apache.flink.util.Preconditions.checkNotNull; import static org.apache.flink.util.Preconditions.checkState; interface RecoveredChannelStateHandler<Info, Context> extends AutoCloseable { @@ -75,7 +85,7 @@ interface RecoveredChannelStateHandler<Info, Context> extends AutoCloseable { /** * Abstract base for all input-channel recovery handlers. Holds the channel mapping logic shared by - * both variants (no-filtering, filtering). + * all three variants (no-spilling, spilling-no-filtering, spilling-with-filtering). * * <p>Subclasses implement {@link #recover} according to their specific recovery mode and override * {@link #closeInternal()} to release mode-specific resources. @@ -143,6 +153,15 @@ abstract class AbstractInputChannelRecoveredStateHandler return new BufferWithContext<>(wrap(buffer), buffer); } + /** + * Returns the {@link FetchedChannelState} produced during spilling, or {@code null} if spilling + * was not active (i.e., {@link NoSpillingHandler}). + */ + @Nullable + FetchedChannelState getProducedChannelState() { + return null; + } + @Override public void close() throws IOException { closeInternal(); @@ -214,6 +233,266 @@ class NoSpillingHandler extends AbstractInputChannelRecoveredStateHandler { } } +/** + * Intermediate abstract base for the two spilling variants. Owns the on-disk spill format end to + * end: a single reusable {@link DataOutputSerializer} accumulates one channel's segment, the + * segment header is backfilled with the body length at seal time, and sealed segments are flushed + * to the current file stream with 64 MB-bounded rotation. + * + * <h3>Disk format</h3> + * + * <pre> + * [ 4B BE int: gate idx ] segment header: written once per channel segment + * [ 4B BE int: channel idx ] + * [ 4B BE int: buffer length ] segment body byte count (backfilled at segment seal) + * [ 4B BE int: record length ] repeated for every record in this segment + * [ N bytes: serialized record ] + * [ 4B BE int: gate idx ] next segment header (channel switch or post-rotation) + * ... + * </pre> + * + * <p>The body byte count is only known after the whole segment is written, so each segment is first + * accumulated in {@link #segmentSerializer} (header written at open with a zero placeholder) and + * {@link DataOutputSerializer#writeIntUnsafe} backfills the length at seal. A segment is one + * uninterrupted run of records for a single channel; file rotation happens only after a segment is + * fully sealed, so a segment never crosses a file boundary. + */ +abstract class AbstractSpillingHandler extends AbstractInputChannelRecoveredStateHandler { + + /** Byte offset of the {@code bufferLength} field within a segment's header. */ + static final int BUFFER_LENGTH_HEADER_OFFSET = 2 * Integer.BYTES; + + /** Total size of the segment header in bytes: gateIdx + channelIdx + bufferLength. */ + static final int SEGMENT_HEADER_BYTES = 3 * Integer.BYTES; + + final String[] spillTmpDirectories; + + public static final long DEFAULT_SPILL_FILE_SIZE_BYTES = 64L * 1024 * 1024; + + public static final int DEFAULT_MAX_SEGMENT_SIZE_BYTES = 1024 * 1024; + + /** Soft per-file size bound that triggers rotation between segments. */ + private final long maxFileSizeBytes; + + /** Soft per-segment size bound that triggers a seal, keeping heap use bounded. */ + private final int maxSegmentSizeBytes; + + /** + * Accumulates the current segment: the header followed by the body, which is either + * length-prefixed filtered records or verbatim pass-through bytes, depending on the subclass. + * Reused across segments via {@code clear()}. + */ + private final DataOutputSerializer segmentSerializer = new DataOutputSerializer(256); + + /** + * Spill files written so far, in order. The {@link FetchedChannelState} handoff is built from + * this list once writing is sealed; an empty list means the handler never spilled any bytes, so + * it produces no state. + */ + private final List<Path> files = new ArrayList<>(); + + /** + * Unique directory for this handler's spill files; created lazily when the first file opens. + */ + private final Path baseDir; + + /** + * Output stream to the current spill file; tracks the bytes written so far via {@link + * OffsetAwareOutputStream#getLength()} to decide when to rotate. Null before the first segment + * is flushed. + */ + @Nullable private OffsetAwareOutputStream currentStream; + + /** Channel whose segment is currently open; null when no segment is in progress. */ + @Nullable private InputChannelInfo currentChannel; + + @Nullable private FetchedChannelState producedChannelState; + + AbstractSpillingHandler( + InputGate[] inputGates, + InflightDataRescalingDescriptor channelMapping, + String[] spillTmpDirectories, + long maxFileSizeBytes, + int maxSegmentSizeBytes) { + // FLINK-38544 transitional: the base's third ctor arg is removed when the spilling backend + // lands (spilling always implies checkpointing-during-recovery enabled). + super(inputGates, channelMapping, true); + checkArgument( + checkNotNull(spillTmpDirectories).length > 0, + "spillTmpDirectories must not be empty"); + checkArgument( + maxFileSizeBytes > 0, "maxFileSizeBytes must be positive: %s", maxFileSizeBytes); + checkArgument( + maxSegmentSizeBytes > 0, + "maxSegmentSizeBytes must be positive: %s", + maxSegmentSizeBytes); + this.spillTmpDirectories = spillTmpDirectories; + this.maxFileSizeBytes = maxFileSizeBytes; + this.maxSegmentSizeBytes = maxSegmentSizeBytes; + this.baseDir = + Paths.get(spillTmpDirectories[0], "flink-channel-spill-" + UUID.randomUUID()); + } + + /** + * Opens (or switches to) the segment for {@code channelInfo} and returns its buffer for the + * caller to append the body into. The caller must not seal the segment. + */ + DataOutputSerializer segmentSerializerFor(InputChannelInfo channelInfo) throws IOException { + startNewSegmentIfNeeded(channelInfo); + return segmentSerializer; + } + + private void startNewSegmentIfNeeded(InputChannelInfo channelInfo) throws IOException { + // A segment ends on a channel switch, or once it outgrew the heap bound; the disk format + // allows consecutive segments for one channel, so the reader is unaffected either way. + if (channelInfo.equals(currentChannel) + && segmentSerializer.length() <= maxSegmentSizeBytes) { + return; + } + if (currentChannel != null) { + sealCurrentSegment(); + } + segmentSerializer.clear(); + segmentSerializer.writeInt(channelInfo.getGateIdx()); + segmentSerializer.writeInt(channelInfo.getInputChannelIdx()); + segmentSerializer.writeInt(0); // bufferLength placeholder + currentChannel = channelInfo; + } + + /** + * Backfills the body length into the segment header and flushes the whole segment to the file + * stream. Empty segments are dropped without opening a file, so no empty file is created. + */ + private void sealCurrentSegment() throws IOException { + if (currentChannel == null) { + return; + } + currentChannel = null; + int totalBytes = segmentSerializer.length(); + int bodyBytes = totalBytes - SEGMENT_HEADER_BYTES; + if (bodyBytes == 0) { + // The header is written before filtering runs, so a channel whose records are all + // filtered out ends up empty: drop it instead of writing a header-only segment. + return; + } + // Math.toIntExact guards against the unlikely case of a single segment > 2 GB. + segmentSerializer.writeIntUnsafe(Math.toIntExact(bodyBytes), BUFFER_LENGTH_HEADER_OFFSET); + ensureFileOpen(); + currentStream.write(segmentSerializer.getSharedBuffer(), 0, totalBytes); + } + + /** + * Ensures an output stream is ready for the next segment, rotating to a fresh file first if the + * current one reached the size bound. Rotation happens here, between sealed segments, so a + * segment is never split across files. + */ + private void ensureFileOpen() throws IOException { + if (currentStream != null && currentStream.getLength() >= maxFileSizeBytes) { + currentStream.flush(); + currentStream.close(); + currentStream = null; + } + if (currentStream != null) { + return; + } + // create the spill dir on the first file; no-op afterwards + Files.createDirectories(baseDir); + Path filePath = baseDir.resolve("spill-segment-" + files.size() + ".bin"); + // CREATE_NEW fails loud if the file already exists instead of silently overwriting it. + currentStream = + new OffsetAwareOutputStream( + new BufferedOutputStream( + Files.newOutputStream( + filePath, + StandardOpenOption.CREATE_NEW, + StandardOpenOption.WRITE)), + 0L); + files.add(filePath); + } + + @Override + @Nullable + FetchedChannelState getProducedChannelState() { + return producedChannelState; + } + + /** Spill files written so far; empty if this handler never spilled any bytes. */ + @VisibleForTesting + List<Path> peekSpillFilesForTesting() { + return files; + } + + /** + * Seals the open segment and the file stream, then builds the {@link FetchedChannelState} + * handoff from the written files. Produces nothing if no bytes were ever spilled. + */ + @Override + void closeInternal() throws IOException { + if (currentChannel != null) { + sealCurrentSegment(); + } + if (currentStream != null) { + currentStream.flush(); + currentStream.close(); // OffsetAwareOutputStream closes the wrapped stream quietly + currentStream = null; + } + if (files.isEmpty()) { + return; + } + producedChannelState = new FetchedChannelState(files); + // Keep the files alive between close() and drain-reader construction. + producedChannelState.acquire(); + } +} + +/** + * Recovery handler for the case where checkpointing during recovery is enabled but no filtering + * handler is present. Appends recovered buffer bytes verbatim into the current segment. + */ +class SpillingNoFilteringHandler extends AbstractSpillingHandler { + + SpillingNoFilteringHandler( + InputGate[] inputGates, + InflightDataRescalingDescriptor channelMapping, + String[] spillTmpDirectories) { + super( + inputGates, + channelMapping, + spillTmpDirectories, + DEFAULT_SPILL_FILE_SIZE_BYTES, + DEFAULT_MAX_SEGMENT_SIZE_BYTES); + } + + @Override + public void recover( + InputChannelInfo channelInfo, + int oldSubtaskIndex, + BufferWithContext<Buffer> bufferWithContext) + throws IOException, InterruptedException { + Buffer buffer = bufferWithContext.context; + try { + if (buffer.readableBytes() > 0) { + recoverPassThroughToSpill(getMappedChannels(channelInfo).getChannelInfo(), buffer); + } + } finally { + buffer.recycleBuffer(); + } + } + + private void recoverPassThroughToSpill(InputChannelInfo channelInfo, Buffer source) + throws IOException { + // The recovered bytes are already a length-prefixed record sequence, so append them + // verbatim into the segment without re-framing. Writing straight from the backing + // MemorySegment lets it absorb the heap/off-heap distinction, avoiding both a branch on the + // NIO buffer kind and the intermediate copy a direct buffer would otherwise require. + segmentSerializerFor(channelInfo) + .write( + source.getMemorySegment(), + source.getMemorySegmentOffset() + source.getReaderIndex(), + source.readableBytes()); + } +} + /** * Recovery handler for the case where checkpointing during recovery is enabled and a filtering * handler is present. Uses a reusable heap-backed pre-filter buffer (isolated from the Network diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/AbstractSpillingHandlerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/AbstractSpillingHandlerTest.java new file mode 100644 index 00000000000..e44a52dbdb3 --- /dev/null +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/AbstractSpillingHandlerTest.java @@ -0,0 +1,184 @@ +/* + * 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.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.DataInputStream; +import java.io.FileInputStream; +import java.nio.file.Path; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Tests the on-disk spill format produced by {@link AbstractSpillingHandler}: the 12-byte segment + * header with backfilled buffer length, channel switching, file rotation, and empty-segment + * dropping. Records are appended through {@link TestSpillWriter}, mirroring how the filtering and + * pass-through subclasses feed the segment buffer. + */ +class AbstractSpillingHandlerTest { + + @TempDir Path tempDir; + + @Test + void testChannelSwitchProducesTwoSegmentsInOneFile() throws Exception { + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writeRecord(new InputChannelInfo(0, 0), bytes(0xAA, 0xBB), 2); + writer.writeRecord(new InputChannelInfo(0, 1), bytes(0xCC, 0xDD, 0xEE), 3); + FetchedChannelState state = writer.getChannelState(); + + assertThat(state.files()).hasSize(1); + int seg0Body = Integer.BYTES + 2; + int seg1Body = Integer.BYTES + 3; + assertThat(state.files().get(0).toFile().length()) + .isEqualTo( + (long) (AbstractSpillingHandler.SEGMENT_HEADER_BYTES + seg0Body) + + (AbstractSpillingHandler.SEGMENT_HEADER_BYTES + seg1Body)); + } + } + + @Test + void testSameChannelContinuousRecordsMergeIntoOneSegment() throws Exception { + InputChannelInfo ch = new InputChannelInfo(0, 0); + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writeRecord(ch, bytes(1), 1); + writer.writeRecord(ch, bytes(2, 3), 2); + writer.writeRecord(ch, bytes(4, 5, 6), 3); + FetchedChannelState state = writer.getChannelState(); + + long expectedBody = 3L * Integer.BYTES + (1 + 2 + 3); + assertThat(state.files()).hasSize(1); + assertThat(state.files().get(0).toFile().length()) + .isEqualTo(AbstractSpillingHandler.SEGMENT_HEADER_BYTES + expectedBody); + } + } + + @Test + void testSameChannelIsSplitIntoSegmentsWhenSizeBoundIsExceeded() throws Exception { + InputChannelInfo ch = new InputChannelInfo(0, 0); + // Segment bound of 8 bytes: every record outgrows it, so each write starts a new segment + // even though the channel never switches. + try (TestSpillWriter writer = + new TestSpillWriter( + tempDir, AbstractSpillingHandler.DEFAULT_SPILL_FILE_SIZE_BYTES, 8)) { + writer.writeRecord(ch, bytes(1, 2, 3, 4, 5, 6, 7, 8), 8); + writer.writeRecord(ch, bytes(1, 2, 3, 4, 5, 6, 7, 8), 8); + writer.writeRecord(ch, bytes(1, 2, 3, 4, 5, 6, 7, 8), 8); + FetchedChannelState state = writer.getChannelState(); + + // Three headers instead of one: 3 * (header + 4B length prefix + 8B payload). + long segmentBytes = AbstractSpillingHandler.SEGMENT_HEADER_BYTES + Integer.BYTES + 8; + assertThat(state.files()).hasSize(1); + assertThat(state.files().get(0).toFile().length()).isEqualTo(3 * segmentBytes); + } + } + + @Test + void testPassThroughBytesAreWrittenVerbatim() throws Exception { + InputChannelInfo ch = new InputChannelInfo(1, 2); + byte[] data = bytes(0x01, 0x02, 0x03, 0x04, 0x05); + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writePassThrough(ch, data, 0, data.length); + FetchedChannelState state = writer.getChannelState(); + + assertThat(state.files()).hasSize(1); + assertThat(state.files().get(0).toFile().length()) + .isEqualTo(AbstractSpillingHandler.SEGMENT_HEADER_BYTES + data.length); + } + } + + // ------------------------------------------------------------------------------------------- + // Empty segments: opening a segment without writing a body must not create a file + // ------------------------------------------------------------------------------------------- + + @Test + void testOpenSegmentWithoutBodyProducesNoFile() throws Exception { + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.openSegment(new InputChannelInfo(0, 0)); + // No body was ever spilled, so no state (and therefore no file) is produced. + assertThat(writer.getChannelState()).isNull(); + } + } + + @Test + void testEmptySegmentsAroundANonEmptyOneProduceExactlyTheNonEmptyFile() throws Exception { + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.openSegment(new InputChannelInfo(0, 0)); + writer.writeRecord(new InputChannelInfo(0, 1), bytes(7, 8, 9), 3); + writer.openSegment(new InputChannelInfo(0, 2)); + FetchedChannelState state = writer.getChannelState(); + + assertThat(state.files()).hasSize(1); + assertThat(state.files().get(0).toFile().length()) + .isEqualTo(AbstractSpillingHandler.SEGMENT_HEADER_BYTES + Integer.BYTES + 3); + } + } + + // ------------------------------------------------------------------------------------------- + // File rotation + // ------------------------------------------------------------------------------------------- + + @Test + void testRotationProducesOneFilePerSegmentWhenBoundIsTiny() throws Exception { + try (TestSpillWriter writer = + new TestSpillWriter( + tempDir, 1L, AbstractSpillingHandler.DEFAULT_MAX_SEGMENT_SIZE_BYTES)) { + writer.writeRecord(new InputChannelInfo(0, 0), bytes(1), 1); + writer.writeRecord(new InputChannelInfo(0, 1), bytes(2, 3), 2); + writer.writeRecord(new InputChannelInfo(0, 2), bytes(4), 1); + FetchedChannelState state = writer.getChannelState(); + + assertThat(state.files()).hasSize(3); + for (Path file : state.files()) { + assertThat(file.toFile().length()).isGreaterThan(0); + } + } + } + + // ------------------------------------------------------------------------------------------- + // Disk-format verification + // ------------------------------------------------------------------------------------------- + + @Test + void testDiskFormatMatchesSpec() throws Exception { + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writeRecord(new InputChannelInfo(7, 3), bytes(0xAB, 0xCD, 0xEF), 3); + Path file = writer.getChannelState().files().get(0); + try (DataInputStream in = new DataInputStream(new FileInputStream(file.toFile()))) { + assertThat(in.readInt()).isEqualTo(7); // gateIdx + assertThat(in.readInt()).isEqualTo(3); // channelIdx + assertThat(in.readInt()).isEqualTo(Integer.BYTES + 3); // bufferLength + assertThat(in.readInt()).isEqualTo(3); // record length prefix + assertThat(in.read()).isEqualTo(0xAB); + assertThat(in.read()).isEqualTo(0xCD); + assertThat(in.read()).isEqualTo(0xEF); + assertThat(in.read()).isEqualTo(-1); // EOF + } + } + } + + private static byte[] bytes(int... values) { + byte[] out = new byte[values.length]; + for (int i = 0; i < values.length; i++) { + out[i] = (byte) values[i]; + } + return out; + } +} diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/TestSpillWriter.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/TestSpillWriter.java new file mode 100644 index 00000000000..de93ecf4990 --- /dev/null +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/TestSpillWriter.java @@ -0,0 +1,105 @@ +/* + * 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.core.memory.DataOutputSerializer; +import org.apache.flink.runtime.checkpoint.InflightDataRescalingDescriptor; +import org.apache.flink.runtime.io.network.partition.consumer.InputGate; + +import java.io.Closeable; +import java.io.IOException; +import java.nio.file.Path; + +/** + * Test-only helper that produces spill files in the {@link AbstractSpillingHandler} on-disk format, + * exposing the same {@code writeRecord} / {@code writePassThrough} surface the readers and drainers + * are tested against. It drives a minimal concrete {@link AbstractSpillingHandler} so tests need + * not stand up real input gates or run the recovery loop. + */ +final class TestSpillWriter implements Closeable { + + private final FormatHandler handler; + + TestSpillWriter(Path baseDir) { + this( + baseDir, + AbstractSpillingHandler.DEFAULT_SPILL_FILE_SIZE_BYTES, + AbstractSpillingHandler.DEFAULT_MAX_SEGMENT_SIZE_BYTES); + } + + TestSpillWriter(Path baseDir, long maxFileSizeBytes, int maxSegmentSizeBytes) { + this.handler = + new FormatHandler( + new String[] {baseDir.toString()}, maxFileSizeBytes, maxSegmentSizeBytes); + } + + /** Appends one length-prefixed record, mirroring the filtering path. */ + void writeRecord(InputChannelInfo channelInfo, byte[] record, int recordLength) + throws IOException { + DataOutputSerializer segment = handler.segmentSerializerFor(channelInfo); + segment.writeInt(recordLength); + segment.write(record, 0, recordLength); + } + + /** Appends verbatim bytes, mirroring the pass-through path. */ + void writePassThrough(InputChannelInfo channelInfo, byte[] data, int offset, int length) + throws IOException { + handler.segmentSerializerFor(channelInfo).write(data, offset, length); + } + + /** Opens (or switches to) a segment without writing any body, to exercise empty segments. */ + void openSegment(InputChannelInfo channelInfo) throws IOException { + handler.segmentSerializerFor(channelInfo); + } + + /** + * Seals the spilled segments and returns the produced state, already holding one lifecycle + * grant for the caller. Returns {@code null} if nothing was ever written. + */ + FetchedChannelState getChannelState() throws IOException { + handler.close(); + return handler.getProducedChannelState(); + } + + @Override + public void close() throws IOException { + handler.close(); + } + + private static final class FormatHandler extends AbstractSpillingHandler { + + FormatHandler( + String[] spillTmpDirectories, long maxFileSizeBytes, int maxSegmentSizeBytes) { + super( + new InputGate[0], + InflightDataRescalingDescriptor.NO_RESCALE, + spillTmpDirectories, + maxFileSizeBytes, + maxSegmentSizeBytes); + } + + @Override + public void recover( + InputChannelInfo info, + int oldSubtaskIndex, + BufferWithContext<org.apache.flink.runtime.io.network.buffer.Buffer> ctx) { + throw new UnsupportedOperationException("not used in format tests"); + } + } +}
