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 46ad84a6133b0b34fbc9cc622d2124de007286d5 Author: Rui Fan <[email protected]> AuthorDate: Mon Jul 6 02:15:24 2026 +0200 [FLINK-40080][checkpoint] FetchedChannelStateDrainer: incremental drain with atomic snapshot-and-insert Implements RecoveryCheckpointTrigger + Closeable; the constructor takes a FetchedChannelState and the task's List<RecoverableInputChannel> and derives the InputChannelInfo map internally. drain() walks the root reader front to back; per segment the recovery buffer request (requestRecoveryBufferBlocking()) and the disk read happen outside the lock, while delivery (onRecoveredStateBuffer) and commit()/offset advance happen inside it; at the end finishRecoveredBufferDelivery() is called per channel. snapshotAndInsertBarriers(cpId) is atomic under the drainer lock: snapshot the committed position (FetchedChannelStateSnapshot, ref-counted acquire) and insert a RecoveryCheckpointBarrier into every in-recovery channel; returns the snapshot's reader (caller-owned). Nothing constructs a drainer in production yet and the handler factory is unswitched, so runtime behavior is unchanged. --- .../channel/FetchedChannelStateDrainer.java | 167 +++++- .../channel/SequentialChannelStateReader.java | 2 - .../channel/SequentialChannelStateReaderImpl.java | 76 +-- .../channel/FetchedChannelStateDrainerTest.java | 663 +++++++++++++++++++++ 4 files changed, 834 insertions(+), 74 deletions(-) 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 23f33bca40e..0fe18769c62 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 @@ -18,71 +18,182 @@ 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.runtime.io.network.buffer.Buffer; import org.apache.flink.runtime.io.network.partition.consumer.RecoverableInputChannel; import java.io.Closeable; import java.io.IOException; +import java.io.InputStream; +import java.util.HashMap; import java.util.List; +import java.util.Map; +import java.util.Optional; +import static org.apache.flink.util.Preconditions.checkArgument; import static org.apache.flink.util.Preconditions.checkNotNull; /** - * Drains a {@link FetchedChannelState} into the physical recovered channels and inserts {@link - * RecoveryCheckpointBarrier}s when a checkpoint fires during recovery. + * Drains a {@link FetchedChannelState} into recovered-buffer queues and snapshots remaining + * segments when a checkpoint fires during recovery. * - * <p>FLINK-38544 transitional in-memory implementation: the in-memory recovery backend already - * pushed every recovered buffer into the physical channels' own queues at conversion time (via - * {@code requestPartitions(true)}), so draining is just appending the end-of-recovered-state - * sentinel to each channel, and there is never an undrained residue to snapshot — inserting the - * barrier into the in-recovery channels is enough. The disk-based drainer of the spilling backend - * replaces this, reading segments off spill files and returning a reader over the undrained slice. + * <p>The drainer lock pairs channel delivery with reader-cursor advancement and also protects + * snapshot creation plus barrier insertion. Disk reads and buffer allocation stay outside that + * lock. */ @Internal public final class FetchedChannelStateDrainer implements RecoveryCheckpointTrigger, Closeable { + private final FetchedChannelStateReader rootReader; + + private final ResolvedChannels channels; + + private final Object lock = new Object(); private final FetchedChannelState channelState; - private final List<RecoverableInputChannel> channels; + /** + * Set under {@link #lock} once {@link #drain()} has consumed every segment. After that the + * {@link #rootReader} is closed by {@link #close()}, so a later {@link + * #snapshotAndInsertBarriers} must not derive from it; it returns an empty reader instead. + * Guarded by the lock so the check is atomic with barrier insertion. + */ + private boolean drainFinished; public FetchedChannelStateDrainer( FetchedChannelState channelState, List<RecoverableInputChannel> channels) { - this.channelState = checkNotNull(channelState); - this.channels = checkNotNull(channels); + this.channelState = channelState; + this.rootReader = checkNotNull(channelState).reader(); + this.channels = new ResolvedChannels(channels); + } + + private static final class ResolvedChannels { + final List<RecoverableInputChannel> allChannels; + final Map<InputChannelInfo, RecoverableInputChannel> channelByInfo; + + ResolvedChannels(List<RecoverableInputChannel> all) { + this.allChannels = all; + Map<InputChannelInfo, RecoverableInputChannel> byInfo = new HashMap<>(); + for (RecoverableInputChannel ch : all) { + byInfo.put(ch.getChannelInfo(), ch); + } + this.channelByInfo = byInfo; + } } /** - * Appends the end-of-recovered-state sentinel to every converted channel. Each channel first - * waits for its upstream to be ready, so this must run on the channelIOExecutor rather than - * block the mailbox thread. Only once the sentinel is in place can the consume path flip the - * channel out of recovery, which guarantees live data is never polled before the upstream - * connection exists. + * Drains all segments from the spill file into the corresponding recovery buffer queues. Each + * segment is split into chunks of at most {@code memorySegmentSize} bytes; a full chunk is + * delivered under the drainer lock paired with a segment commit. After all segments are + * drained, every channel's {@link RecoverableInputChannel#finishRecoveredBufferDelivery()} is + * called. + * + * <p>Disk reads and buffer allocations happen outside the lock; only the "deliver + commit" + * pair is locked to guarantee atomicity with snapshot. */ public void drain() throws IOException, InterruptedException { channelState.release(); - for (RecoverableInputChannel channel : channels) { - channel.finishRecoveredBufferDelivery(); + Optional<SpillSegment> next; + while ((next = rootReader.advanceAndGetNextSegment()).isPresent()) { + SpillSegment seg = next.get(); + RecoverableInputChannel ch = channels.channelByInfo.get(seg.channelInfo()); + if (ch == null) { + throw new IllegalStateException( + "Drain: no physical channel found for " + seg.channelInfo()); + } + drainSegment(seg, ch); + } + + // Mark drain done before rootReader is closed, so a concurrent snapshot returns empty + // rather than deriving from the soon-to-be-closed rootReader. Under the lock to stay atomic + // with snapshotAndInsertBarriers' check. + synchronized (lock) { + drainFinished = true; + } + for (RecoverableInputChannel ch : channels.allChannels) { + ch.finishRecoveredBufferDelivery(); } } /** - * Inserts a {@link RecoveryCheckpointBarrier} into every channel that is still in recovery, so - * that {@code checkpointStarted}'s in-recovery branch can persist exactly the pre-barrier - * recovered data. There is no snapshot side for the in-memory backend: everything a checkpoint - * must persist is already inside the channels' queues, so the snapshot is inherently empty. + * Drains one segment into the given channel, delivering a buffer under the lock once it is full + * or the segment is exhausted. + * + * <p>{@link RecoverableInputChannel#onRecoveredStateBuffer} takes ownership even when it + * throws, so the reference is dropped before the call; whatever is still held is recycled on + * error. + */ + private void drainSegment(SpillSegment seg, RecoverableInputChannel ch) + throws IOException, InterruptedException { + InputStream in = seg.bodyStream(); + int remaining = seg.length(); + Buffer buf = null; + try { + while (remaining > 0) { + if (buf == null) { + buf = ch.requestRecoveryBufferBlocking(); + } + remaining -= + fill(buf, in, Math.min(buf.getMaxCapacity() - buf.getSize(), remaining)); + if (buf.getSize() == buf.getMaxCapacity() || remaining == 0) { + Buffer delivered = buf; + buf = null; + synchronized (lock) { + ch.onRecoveredStateBuffer(delivered); + seg.commit(); + } + } + } + } catch (Throwable t) { + if (buf != null) { + buf.recycleBuffer(); + } + throw t; + } + } + + /** + * Writes up to {@code toRead} bytes from {@code in} into {@code buf} and returns how many were + * written. Does not close or recycle {@code buf}; ownership stays with the caller. + */ + private static int fill(Buffer buf, InputStream in, int toRead) throws IOException { + checkArgument(toRead > 0); + // Do not use try-with-resources: ChannelStateByteBuffer.close() recycles the buffer, + // but the buffer is still owned by the caller here. + ChannelStateByteBuffer view = ChannelStateByteBuffer.wrap(buf); + return view.writeBytes(in, toRead); + } + + /** + * Atomically snapshots the undrained portion of the spill and inserts {@link + * RecoveryCheckpointBarrier}s into all in-recovery channels. Returns an independent reader over + * the remaining segments for replay into the checkpoint stream; the caller owns and must close + * it. + * + * <p>If the drain has already finished, the root reader is closed and there is nothing left to + * snapshot; an empty reader is returned so the caller's normal flow handles it uniformly. */ @Override public FetchedChannelStateSnapshot snapshotAndInsertBarriers(long checkpointId) throws IOException { - for (RecoverableInputChannel channel : channels) { - channel.insertRecoveryCheckpointBarrierIfInRecovery(checkpointId); + + // Barrier insertion and snapshot must occur within the same critical section so that the + // snapshot's committed position reflects exactly the drain position at the moment barriers + // were inserted, with no window for the drain thread to advance between. + synchronized (lock) { + for (RecoverableInputChannel ch : channels.allChannels) { + ch.insertRecoveryCheckpointBarrierIfInRecovery(checkpointId); + } + if (drainFinished) { + // Drain consumed everything and rootReader is (being) closed; nothing left to + // snapshot. Return an empty snapshot so the caller's normal flow handles it. + return FetchedChannelState.emptySnapshot(); + } + return rootReader.snapshot(); } - // 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 public void close() throws IOException { - channelState.close(); + rootReader.close(); } } diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReader.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReader.java index f6b8decf09c..88296c517b0 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReader.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReader.java @@ -34,8 +34,6 @@ public interface SequentialChannelStateReader extends AutoCloseable { * * @param inputGates The input gates to recover state for. * @param filterContext The filter context containing input configs and rescaling info. - * @return the fetched recovered channel state if there is any state to recover, otherwise - * {@link Optional#empty()}. */ Optional<FetchedChannelState> readInputData( InputGate[] inputGates, RecordFilterContext filterContext) diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java index 784aa7c1c0d..9284c226a59 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java @@ -33,7 +33,6 @@ import org.apache.flink.streaming.runtime.io.recovery.RecordFilterContext; import java.io.Closeable; import java.io.IOException; import java.util.Collection; -import java.util.Collections; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -72,45 +71,38 @@ public class SequentialChannelStateReaderImpl implements SequentialChannelStateR ? ChannelStateFilteringHandler.createFromContext(filterContext, inputGates) : null; - try (ChannelStateFilteringHandler ignored = filteringHandler; - AbstractInputChannelRecoveredStateHandler stateHandler = - AbstractInputChannelRecoveredStateHandler.create( - inputGates, - taskStateSnapshot.getInputRescalingDescriptor(), - filterContext.isCheckpointingDuringRecoveryEnabled(), - filteringHandler, - filterContext.getMemorySegmentSize(), - filterContext.getTmpDirectories())) { - boolean readAny = - read( - stateHandler, - groupByDelegate( - streamSubtaskStates(), - ChannelStateHelper::extractUnmergedInputHandles)); - readAny |= - read( - stateHandler, - groupByDelegate( - streamSubtaskStates(), - OperatorSubtaskState::getUpstreamOutputBufferState)); - - if (filteringHandler != null) { - checkState( - !filteringHandler.hasPartialData(), - "Not all data has been fully consumed during filtering"); + // Manual close ordering so the produced spill file can be published after + // stateHandler.close() flushes the filter writer. + AbstractInputChannelRecoveredStateHandler stateHandler = + AbstractInputChannelRecoveredStateHandler.create( + inputGates, + taskStateSnapshot.getInputRescalingDescriptor(), + filterContext.isCheckpointingDuringRecoveryEnabled(), + filteringHandler, + filterContext.getMemorySegmentSize(), + filterContext.getTmpDirectories()); + try (ChannelStateFilteringHandler ignored = filteringHandler) { + try (stateHandler) { + read( + stateHandler, + groupByDelegate( + streamSubtaskStates(), + ChannelStateHelper::extractUnmergedInputHandles)); + read( + stateHandler, + groupByDelegate( + streamSubtaskStates(), + OperatorSubtaskState::getUpstreamOutputBufferState)); + + if (filteringHandler != null) { + checkState( + !filteringHandler.hasPartialData(), + "Not all data has been fully consumed during filtering"); + } } - // The container signals "any recovered data was pushed", not "a filter ran": the - // no-checkpointing path pushes recovered buffers directly and must produce nothing - // here, matching the caller's checkState(readInputData(...).isEmpty()). - // - // FLINK-38544 transitional in-memory backend: recovered buffers already live in the - // physical channels' queues, so the returned container is an empty placeholder. - // TODO: FLINK-38544 — return stateHandler.getProducedChannelState() once the handler - // factory selects the Spilling* handlers; as-is their file-backed state would be - // silently dropped here. - return filterContext.isCheckpointingDuringRecoveryEnabled() && readAny - ? Optional.of(new FetchedChannelState(Collections.emptyList())) - : Optional.empty(); + // stateHandler.close() (above) has flushed the filter writer and published the + // produced spill file, so read getProducedChannelState() after the close completes. + return Optional.ofNullable(stateHandler.getProducedChannelState()); } } @@ -130,19 +122,15 @@ public class SequentialChannelStateReaderImpl implements SequentialChannelStateR } } - /** Returns {@code true} if any channel state handle was read. */ - private <Info, Context, Handle extends AbstractChannelStateHandle<Info>> boolean read( + private <Info, Context, Handle extends AbstractChannelStateHandle<Info>> void read( RecoveredChannelStateHandler<Info, Context> stateHandler, Map<StreamStateHandle, List<Handle>> streamStateHandleListMap) throws IOException, InterruptedException { - boolean readAny = false; for (Map.Entry<StreamStateHandle, List<Handle>> delegateAndHandles : streamStateHandleListMap.entrySet()) { readSequentially( delegateAndHandles.getKey(), delegateAndHandles.getValue(), stateHandler); - readAny = true; } - return readAny; } private <Info, Context, Handle extends AbstractChannelStateHandle<Info>> void readSequentially( diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateDrainerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateDrainerTest.java new file mode 100644 index 00000000000..d729bef9e98 --- /dev/null +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateDrainerTest.java @@ -0,0 +1,663 @@ +/* + * 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.MemorySegment; +import org.apache.flink.core.memory.MemorySegmentFactory; +import org.apache.flink.runtime.checkpoint.channel.FetchedChannelStateReader.SpillSegment; +import org.apache.flink.runtime.event.AbstractEvent; +import org.apache.flink.runtime.io.network.api.serialization.EventSerializer; +import org.apache.flink.runtime.io.network.buffer.Buffer; +import org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler; +import org.apache.flink.runtime.io.network.buffer.NetworkBuffer; +import org.apache.flink.runtime.io.network.partition.consumer.RecoverableInputChannel; + +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.channels.FileChannel; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Tests for {@link FetchedChannelStateDrainer}: drain demux, finish ordering, snapshot start + * position, barrier insertion, and edge cases (drain-finished, channel-not-in-recovery). + */ +class FetchedChannelStateDrainerTest { + + @TempDir Path tempDir; + + @Test + void testDrainEndToEnd() throws Exception { + InputChannelInfo cInfo = new InputChannelInfo(0, 0); + FetchedChannelState state = writeRecords(cInfo, payload(1), payload(2), payload(3)); + + RecordingChannel rec = new RecordingChannel(cInfo); + FetchedChannelStateDrainer drainer = newDrainer(state, cInfo, rec); + + drainer.drain(); + drainer.close(); + + // All segment bodies must be delivered as buffer(s); at least 3 non-empty deliveries + // because the segment body contains 3 records but they may be batched into fewer buffers. + assertThat(rec.recovered).isNotEmpty(); + assertThat(rec.finishCalls).isEqualTo(1); + } + + @Test + void testDrainSegmentLargerThanBufferSplitsIntoFullChunksThenPartialTail() throws Exception { + InputChannelInfo cInfo = new InputChannelInfo(0, 0); + + // Buffer capacity deliberately smaller than the segment body so the drainer must fill + // multiple buffers and a final partial tail. 50 bytes over a 16-byte buffer => 16+16+16+2. + int bufferCapacity = 16; + byte[] body = sequentialBytes(50); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + // Pass-through so the segment body equals the verbatim bytes (no length framing). + writer.writePassThrough(cInfo, body, 0, body.length); + state = writer.getChannelState(); + } + + RecordingChannel rec = new RecordingChannel(cInfo, bufferCapacity); + FetchedChannelStateDrainer drainer = newDrainer(state, cInfo, rec); + + drainer.drain(); + drainer.close(); + + // ceil(50 / 16) = 4 buffers delivered. + assertThat(rec.recovered).hasSize(4); + // Every buffer except the last is filled to capacity; the last carries the remainder. + for (int i = 0; i < rec.recovered.size() - 1; i++) { + assertThat(rec.recovered.get(i).getSize()).isEqualTo(bufferCapacity); + } + assertThat(rec.recovered.get(rec.recovered.size() - 1).getSize()) + .isEqualTo(body.length % bufferCapacity); + + // Buffers concatenated in delivery order must reproduce the segment body byte-for-byte. + assertThat(concat(rec.recovered)).isEqualTo(body); + assertThat(rec.finishCalls).isEqualTo(1); + } + + @Test + void testDrainSegmentExactMultipleOfBufferHasNoPartialTail() throws Exception { + InputChannelInfo cInfo = new InputChannelInfo(0, 0); + + // Body length is an exact multiple of the buffer capacity: the segment ends on a buffer + // boundary, so no extra buffer is requested and no empty buffer is delivered. + int bufferCapacity = 16; + byte[] body = sequentialBytes(bufferCapacity * 3); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writePassThrough(cInfo, body, 0, body.length); + state = writer.getChannelState(); + } + + RecordingChannel rec = new RecordingChannel(cInfo, bufferCapacity); + FetchedChannelStateDrainer drainer = newDrainer(state, cInfo, rec); + + drainer.drain(); + drainer.close(); + + // Exactly 3 full buffers, no trailing empty buffer. + assertThat(rec.recovered).hasSize(3); + for (Buffer b : rec.recovered) { + assertThat(b.getSize()).isEqualTo(bufferCapacity); + } + assertThat(concat(rec.recovered)).isEqualTo(body); + assertThat(rec.finishCalls).isEqualTo(1); + } + + @Test + void testDrainRecyclesInFlightBufferWhenBodyReadFails() throws Exception { + InputChannelInfo cInfo = new InputChannelInfo(0, 0); + + int bufferCapacity = 16; + byte[] body = sequentialBytes(bufferCapacity * 4); + + FetchedChannelState state; + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + writer.writePassThrough(cInfo, body, 0, body.length); + state = writer.getChannelState(); + } + + // Cut the file short so the body read fails while the drainer still owns a partially + // filled buffer: 40 of the 64 body bytes survive => 2 buffers delivered, the third one + // holds 8 bytes when the EOF hits and must be recycled rather than leaked. + try (FileChannel file = FileChannel.open(state.files().get(0), StandardOpenOption.WRITE)) { + file.truncate(file.size() - 24); + } + + RecordingChannel rec = new RecordingChannel(cInfo, bufferCapacity); + FetchedChannelStateDrainer drainer = newDrainer(state, cInfo, rec); + + assertThatThrownBy(drainer::drain).isInstanceOf(EOFException.class); + drainer.close(); + + assertThat(rec.recovered).hasSize(2); + assertThat(rec.requested).isEqualTo(3); + assertThat(rec.recycled).isEqualTo(1); + } + + @Test + void testDrainDemuxByChannelInfo() 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, payload(11), payload(11).length); + writer.writeRecord(c1, payload(22), payload(22).length); + writer.writeRecord(c0, payload(33), payload(33).length); + writer.writeRecord(c1, payload(44), payload(44).length); + state = writer.getChannelState(); + } + + RecordingChannel chan0 = new RecordingChannel(c0); + RecordingChannel chan1 = new RecordingChannel(c1); + FetchedChannelStateDrainer drainer = newDrainer(state, c0, chan0, c1, chan1); + + drainer.drain(); + drainer.close(); + + // Each channel must receive some data buffers + assertThat(chan0.recovered).isNotEmpty(); + assertThat(chan1.recovered).isNotEmpty(); + // Both channels must have finish called + assertThat(chan0.finishCalls).isEqualTo(1); + assertThat(chan1.finishCalls).isEqualTo(1); + } + + @Test + void testDrainCallsFinishAfterAllBufferDeliveries() 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, payload(1), payload(1).length); + writer.writeRecord(c1, payload(2), payload(2).length); + state = writer.getChannelState(); + } + + int[] seq = {0}; + RecordingChannel chan0 = new RecordingChannel(c0, seq); + RecordingChannel chan1 = new RecordingChannel(c1, seq); + FetchedChannelStateDrainer drainer = newDrainer(state, c0, chan0, c1, chan1); + + drainer.drain(); + drainer.close(); + + int maxDataSeq = Math.max(chan0.maxDataSeq, chan1.maxDataSeq); + int minFinishSeq = Math.min(chan0.finishSeq, chan1.finishSeq); + assertThat(maxDataSeq).isLessThan(minFinishSeq); + } + + @Test + void testSnapshotCoversAllSegmentsBeforeDrain() throws Exception { + InputChannelInfo cInfo = new InputChannelInfo(0, 0); + FetchedChannelState state = writeRecords(cInfo, payload(5), payload(6)); + + RecordingChannel chan = new RecordingChannel(cInfo); + FetchedChannelStateDrainer drainer = newDrainer(state, cInfo, chan); + + long cpId = 42L; + FetchedChannelStateReader snap = drainer.snapshotAndInsertBarriers(cpId).reader(); + + // Snapshot must cover all segments (at least 1 segment for cInfo). + // The sequential reader requires each segment body to be fully consumed before advancing, + // mirroring the real consumer (ChannelStateCheckpointWriter#writeInputFromSpill). + int count = 0; + Optional<SpillSegment> next; + while ((next = snap.advanceAndGetNextSegment()).isPresent()) { + drainBody(next.get().bodyStream()); + count++; + } + snap.close(); + assertThat(count).isGreaterThan(0); + drainer.close(); + } + + @Test + void testSnapshotInsertsBarrierPerChannel() 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, payload(1), payload(1).length); + writer.writeRecord(c1, payload(2), payload(2).length); + state = writer.getChannelState(); + } + + RecordingChannel chan0 = new RecordingChannel(c0); + RecordingChannel chan1 = new RecordingChannel(c1); + FetchedChannelStateDrainer drainer = newDrainer(state, c0, chan0, c1, chan1); + + long cpId = 7L; + FetchedChannelStateReader snap = drainer.snapshotAndInsertBarriers(cpId).reader(); + snap.close(); + + assertThat(chan0.recovered).hasSize(1); + assertThat(chan1.recovered).hasSize(1); + assertThat(extractRecoveryBarrierCheckpointId(chan0.recovered.get(0))).isEqualTo(cpId); + assertThat(extractRecoveryBarrierCheckpointId(chan1.recovered.get(0))).isEqualTo(cpId); + drainer.close(); + } + + @Test + void testSnapshotInsertsBarrierWhenChannelInRecoveryEvenIfDiskSliceEmpty() throws Exception { + InputChannelInfo cInfo = new InputChannelInfo(0, 0); + FetchedChannelState state = writeRecords(cInfo, payload(1)); + + RecordingChannel chan = new RecordingChannel(cInfo); + FetchedChannelStateDrainer drainer = newDrainer(state, cInfo, chan); + + drainer.drain(); + // Drain finished; simulate the channel still in recovery + chan.inRecovery = true; + int recoveredBefore = chan.recovered.size(); + + long cpId = 6L; + FetchedChannelStateReader snap = drainer.snapshotAndInsertBarriers(cpId).reader(); + assertThat(snap.advanceAndGetNextSegment()).isEmpty(); + snap.close(); + + // Barrier must be inserted even though disk slice is empty + assertThat(chan.recovered).hasSize(recoveredBefore + 1); + assertThat(extractRecoveryBarrierCheckpointId(chan.recovered.get(recoveredBefore))) + .isEqualTo(cpId); + drainer.close(); + } + + @Test + void testSnapshotInsertsBarrierOnlyForChannelsStillInRecovery() 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, payload(1), payload(1).length); + state = writer.getChannelState(); + } + + RecordingChannel chan0 = new RecordingChannel(c0); + RecordingChannel chan1 = new RecordingChannel(c1); + chan1.inRecovery = false; + + FetchedChannelStateDrainer drainer = newDrainer(state, c0, chan0, c1, chan1); + + long cpId = 11L; + FetchedChannelStateReader snap = drainer.snapshotAndInsertBarriers(cpId).reader(); + snap.close(); + + assertThat(chan0.recovered).hasSize(1); + assertThat(extractRecoveryBarrierCheckpointId(chan0.recovered.get(0))).isEqualTo(cpId); + assertThat(chan1.recovered).isEmpty(); + drainer.close(); + } + + @Test + void testSnapshotReturnsEmptyWhenDrainFinishedAndNotInRecovery() throws Exception { + InputChannelInfo cInfo = new InputChannelInfo(0, 0); + FetchedChannelState state = writeRecords(cInfo, payload(1), payload(2)); + + RecordingChannel chan = new RecordingChannel(cInfo); + FetchedChannelStateDrainer drainer = newDrainer(state, cInfo, chan); + + drainer.drain(); + chan.inRecovery = false; + int recoveredBefore = chan.recovered.size(); + + FetchedChannelStateReader snap = drainer.snapshotAndInsertBarriers(99L).reader(); + assertThat(snap.advanceAndGetNextSegment()).isEmpty(); + snap.close(); + + // No barrier added since channel left recovery + assertThat(chan.recovered).hasSize(recoveredBefore); + drainer.close(); + } + + @Test + void testSnapshotAfterDrainerClosedReturnsEmptyWithoutTouchingClosedRootReader() + throws Exception { + // Mirrors production order: drain() then close() (which closes the root reader) run before + // a + // late checkpoint fires snapshotAndInsertBarriers. The drain-finished flag must + // short-circuit + // so the closed root reader is never snapshotted. + InputChannelInfo cInfo = new InputChannelInfo(0, 0); + FetchedChannelState state = writeRecords(cInfo, payload(1), payload(2)); + + RecordingChannel chan = new RecordingChannel(cInfo); + FetchedChannelStateDrainer drainer = newDrainer(state, cInfo, chan); + + drainer.drain(); + drainer.close(); + + FetchedChannelStateReader snap = drainer.snapshotAndInsertBarriers(99L).reader(); + assertThat(snap.advanceAndGetNextSegment()).isEmpty(); + snap.close(); + } + + @Test + void testDrainOnExecutorThreadDeliversAndFinishes() throws Exception { + InputChannelInfo cInfo = new InputChannelInfo(0, 0); + FetchedChannelState state = writeRecords(cInfo, payload(1)); + + CapturingChannel chan = new CapturingChannel(cInfo); + FetchedChannelStateDrainer drainer = newDrainer(state, cInfo, chan); + + ExecutorService channelIOExecutor = Executors.newSingleThreadExecutor(); + try { + CompletableFuture<Void> done = new CompletableFuture<>(); + channelIOExecutor.execute( + () -> { + try { + drainer.drain(); + done.complete(null); + } catch (Throwable t) { + done.completeExceptionally(t); + } finally { + try { + drainer.close(); + } catch (IOException ignore) { + } + } + }); + + done.get(5, TimeUnit.SECONDS); + assertThat(chan.dataDeliveries).isGreaterThan(0); + assertThat(chan.finishCalled).isTrue(); + } finally { + channelIOExecutor.shutdownNow(); + assertThat(channelIOExecutor.awaitTermination(5, TimeUnit.SECONDS)).isTrue(); + } + } + + @Test + void testDrainOnExecutorThreadBubblesDeliveryException() throws Exception { + InputChannelInfo cInfo = new InputChannelInfo(0, 0); + FetchedChannelState state = writeRecords(cInfo, payload(1)); + + RecoverableInputChannel chan = + new RecoverableInputChannel() { + @Override + public InputChannelInfo getChannelInfo() { + return cInfo; + } + + @Override + public void onRecoveredStateBuffer(Buffer buffer) { + throw new RuntimeException("boom"); + } + + @Override + public void finishRecoveredBufferDelivery() {} + + @Override + public void insertRecoveryCheckpointBarrierIfInRecovery(long checkpointId) { + throw new RuntimeException("boom"); + } + + @Override + public Buffer requestRecoveryBufferBlocking() { + MemorySegment seg = MemorySegmentFactory.allocateUnpooledSegment(64); + return new NetworkBuffer(seg, FreeingBufferRecycler.INSTANCE); + } + + @Override + public void onRecoveredStateConsumed() {} + }; + + FetchedChannelStateDrainer drainer = newDrainer(state, cInfo, chan); + + CountDownLatch handlerCalled = new CountDownLatch(1); + AtomicReference<Throwable> captured = new AtomicReference<>(); + ExecutorService channelIOExecutor = Executors.newSingleThreadExecutor(); + try { + channelIOExecutor.execute( + () -> { + try { + drainer.drain(); + } catch (Throwable t) { + captured.set(t); + handlerCalled.countDown(); + } finally { + try { + drainer.close(); + } catch (IOException ignore) { + } + } + }); + + assertThat(handlerCalled.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(captured.get()).isInstanceOf(RuntimeException.class); + assertThat(captured.get().getMessage()).isEqualTo("boom"); + } finally { + channelIOExecutor.shutdownNow(); + assertThat(channelIOExecutor.awaitTermination(5, TimeUnit.SECONDS)).isTrue(); + } + } + + // ------------------------------------------------------------------------------------------- + // Helpers + // ------------------------------------------------------------------------------------------- + + private FetchedChannelState writeRecords(InputChannelInfo ch, byte[]... payloads) + throws IOException { + try (TestSpillWriter writer = new TestSpillWriter(tempDir)) { + for (byte[] p : payloads) { + writer.writeRecord(ch, p, p.length); + } + return writer.getChannelState(); + } + } + + private FetchedChannelStateDrainer newDrainer( + FetchedChannelState state, Object... infoChannelPairs) { + List<RecoverableInputChannel> all = new ArrayList<>(); + for (int i = 0; i < infoChannelPairs.length; i += 2) { + all.add((RecoverableInputChannel) infoChannelPairs[i + 1]); + } + return new FetchedChannelStateDrainer(state, all); + } + + private static long extractRecoveryBarrierCheckpointId(Buffer buffer) throws IOException { + AbstractEvent event = + EventSerializer.fromBuffer( + buffer, RecoveryCheckpointBarrier.class.getClassLoader()); + buffer.setReaderIndex(0); + assertThat(event).isInstanceOf(RecoveryCheckpointBarrier.class); + return ((RecoveryCheckpointBarrier) event).getCheckpointId(); + } + + private static byte[] payload(int id) { + return new byte[] {(byte) (id & 0xff), (byte) ((id >> 8) & 0xff), (byte) 0xAB, (byte) 0xCD}; + } + + /** Builds {@code n} bytes whose values count up modulo 256, so order mismatches are visible. */ + private static byte[] sequentialBytes(int n) { + byte[] out = new byte[n]; + for (int i = 0; i < n; i++) { + out[i] = (byte) i; + } + return out; + } + + /** Concatenates the readable bytes of the given buffers in order. */ + private static byte[] concat(List<Buffer> buffers) { + java.io.ByteArrayOutputStream out = new java.io.ByteArrayOutputStream(); + for (Buffer b : buffers) { + java.nio.ByteBuffer nio = b.getNioBufferReadable(); + byte[] chunk = new byte[nio.remaining()]; + nio.get(chunk); + out.write(chunk, 0, chunk.length); + } + return out.toByteArray(); + } + + /** Fully consumes a segment body so the sequential reader may advance to the next segment. */ + private static void drainBody(InputStream body) throws IOException { + byte[] buf = new byte[256]; + while (body.read(buf) != -1) { + // discard + } + } + + // ------------------------------------------------------------------------------------------- + // RecordingChannel stub + // ------------------------------------------------------------------------------------------- + + private static final int DEFAULT_RECOVERY_BUFFER_CAPACITY = 4096; + + private static final class RecordingChannel implements RecoverableInputChannel { + private final InputChannelInfo channelInfo; + final List<Buffer> recovered = new ArrayList<>(); + int finishCalls = 0; + private final int[] sequence; + private final int bufferCapacity; + int requested = 0; + int recycled = 0; + int maxDataSeq = Integer.MIN_VALUE; + int finishSeq = -1; + boolean inRecovery = true; + + RecordingChannel(InputChannelInfo channelInfo) { + this(channelInfo, null, DEFAULT_RECOVERY_BUFFER_CAPACITY); + } + + RecordingChannel(InputChannelInfo channelInfo, int[] sharedSequence) { + this(channelInfo, sharedSequence, DEFAULT_RECOVERY_BUFFER_CAPACITY); + } + + RecordingChannel(InputChannelInfo channelInfo, int bufferCapacity) { + this(channelInfo, null, bufferCapacity); + } + + RecordingChannel(InputChannelInfo channelInfo, int[] sharedSequence, int bufferCapacity) { + this.channelInfo = channelInfo; + this.sequence = sharedSequence; + this.bufferCapacity = bufferCapacity; + } + + @Override + public InputChannelInfo getChannelInfo() { + return channelInfo; + } + + @Override + public void onRecoveredStateBuffer(Buffer buffer) { + recovered.add(buffer); + if (sequence != null) { + maxDataSeq = Math.max(maxDataSeq, ++sequence[0]); + } + } + + @Override + public void finishRecoveredBufferDelivery() { + finishCalls++; + if (sequence != null) { + finishSeq = ++sequence[0]; + } + } + + @Override + public void insertRecoveryCheckpointBarrierIfInRecovery(long checkpointId) + throws IOException { + if (inRecovery) { + recovered.add( + EventSerializer.toBuffer( + new RecoveryCheckpointBarrier(checkpointId), false)); + } + } + + @Override + public Buffer requestRecoveryBufferBlocking() { + requested++; + MemorySegment seg = MemorySegmentFactory.allocateUnpooledSegment(bufferCapacity); + return new NetworkBuffer( + seg, + memorySegment -> { + recycled++; + FreeingBufferRecycler.INSTANCE.recycle(memorySegment); + }); + } + + @Override + public void onRecoveredStateConsumed() {} + } + + /** Counts data deliveries and finish for the executor-thread drain tests. */ + private static final class CapturingChannel implements RecoverableInputChannel { + private final InputChannelInfo channelInfo; + int dataDeliveries = 0; + boolean finishCalled = false; + + CapturingChannel(InputChannelInfo channelInfo) { + this.channelInfo = channelInfo; + } + + @Override + public InputChannelInfo getChannelInfo() { + return channelInfo; + } + + @Override + public void onRecoveredStateBuffer(Buffer buffer) { + if (buffer.isBuffer()) { + dataDeliveries++; + } + } + + @Override + public void finishRecoveredBufferDelivery() { + finishCalled = true; + } + + @Override + public void insertRecoveryCheckpointBarrierIfInRecovery(long checkpointId) {} + + @Override + public Buffer requestRecoveryBufferBlocking() { + MemorySegment seg = MemorySegmentFactory.allocateUnpooledSegment(64); + return new NetworkBuffer(seg, FreeingBufferRecycler.INSTANCE); + } + + @Override + public void onRecoveredStateConsumed() {} + } +}
