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 5a2151f091ddc8cfbb5c68ad000735c493b5aa8e Author: Rui Fan <[email protected]> AuthorDate: Mon Jul 6 00:25:16 2026 +0200 [FLINK-39521][network] LocalInputChannel: push-based recovery state Teach LocalInputChannel to be created directly in a recovery state and to receive recovered buffers through the RecoverableInputChannel interface, instead of migrating them via the constructor: - Constructor: drop ArrayDeque<Buffer> initialRecoveredBuffers, add networkBuffersPerChannel and needsRecovery. With needsRecovery=false the stateConsumedFuture is pre-completed, inRecovery stays false and no recovery BufferManager exists -- behavior is byte-for-byte unchanged. - recoveredBuffers becomes a Deque<Buffer> whose monitor also guards inRecovery and recoverySequenceNumber (starts at Integer.MIN_VALUE). - RecoverableInputChannel implementation: push append with released-state recycling; finishRecoveredBufferDelivery() waits for upstreamReady and appends the EndOfFetchedChannelStateEvent sentinel; insertRecoveryCheckpointBarrierIfInRecovery(); requestRecoveryBufferBlocking() lends exclusive buffers via the recovery-only BufferManager (setup() requests them) -- no production caller until the drainer PR; onRecoveredStateConsumed() flips recovery off and completes the stateConsumedFuture. - getNextBuffer() is rewritten under a single inRecovery predicate: recovered buffer first, then a pending priority event pulled from the subpartition view, otherwise empty (live data stays hidden while in recovery). wrapRecoveredBufferAsAvailability() materializes FileRegionBuffer/CompositeBuffer, assigns recovery sequence numbers and computes the next data type via peekNextDataType. - upstreamReady completes on subpartition-view creation; completeUpstreamReadyForTest() for tests. - checkpointStarted() splits into mutually exclusive in-recovery/normal branches. In recovery it only starts persisting (the persist window stays open) with the buffers collected by collectPreRecoveryBarrier, which walks the queue to the matching RecoveryCheckpointBarrier, retains pre-barrier data buffers and removes the sentinel; a missing barrier is an IOException wrapped as CheckpointException(CHECKPOINT_DECLINED). Mechanical constructor-caller adaptations ride along (UnknownInputChannel, LocalRecoveredInputChannel, InputChannelBuilder, benchmark factory), all passing needsRecovery=false at this stage. Includes a port of the FLINK-40016 double-persist regression test to the push-based recovery API. --- .../partition/consumer/LocalInputChannel.java | 576 ++++++++++++++++----- .../consumer/LocalRecoveredInputChannel.java | 5 +- .../partition/consumer/UnknownInputChannel.java | 3 +- .../partition/consumer/InputChannelBuilder.java | 18 +- .../partition/consumer/LocalInputChannelTest.java | 470 ++++++++++++++--- .../benchmark/SingleInputGateBenchmarkFactory.java | 3 +- 6 files changed, 890 insertions(+), 185 deletions(-) diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/LocalInputChannel.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/LocalInputChannel.java index 33338000ec8..9a7aa4b96df 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/LocalInputChannel.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/LocalInputChannel.java @@ -21,11 +21,15 @@ package org.apache.flink.runtime.io.network.partition.consumer; import org.apache.flink.annotation.VisibleForTesting; import org.apache.flink.metrics.Counter; import org.apache.flink.runtime.checkpoint.CheckpointException; +import org.apache.flink.runtime.checkpoint.CheckpointFailureReason; import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter; +import org.apache.flink.runtime.checkpoint.channel.RecoveryCheckpointBarrier; +import org.apache.flink.runtime.event.AbstractEvent; import org.apache.flink.runtime.event.TaskEvent; import org.apache.flink.runtime.execution.CancelTaskException; import org.apache.flink.runtime.io.network.TaskEventPublisher; import org.apache.flink.runtime.io.network.api.CheckpointBarrier; +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.CompositeBuffer; import org.apache.flink.runtime.io.network.buffer.FileRegionBuffer; @@ -38,28 +42,32 @@ import org.apache.flink.runtime.io.network.partition.ResultPartitionManager; import org.apache.flink.runtime.io.network.partition.ResultSubpartition.BufferAndBacklog; import org.apache.flink.runtime.io.network.partition.ResultSubpartitionIndexSet; import org.apache.flink.runtime.io.network.partition.ResultSubpartitionView; -import org.apache.flink.util.concurrent.FutureUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import javax.annotation.Nullable; +import javax.annotation.concurrent.GuardedBy; import java.io.IOException; import java.util.ArrayDeque; import java.util.ArrayList; +import java.util.Collections; import java.util.Deque; +import java.util.Iterator; import java.util.List; import java.util.Optional; import java.util.Timer; import java.util.TimerTask; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; import static org.apache.flink.util.Preconditions.checkNotNull; import static org.apache.flink.util.Preconditions.checkState; /** An input channel, which requests a local subpartition. */ -public class LocalInputChannel extends InputChannel implements BufferAvailabilityListener { +public class LocalInputChannel extends InputChannel + implements BufferAvailabilityListener, RecoverableInputChannel { private static final Logger LOG = LoggerFactory.getLogger(LocalInputChannel.class); @@ -83,19 +91,53 @@ public class LocalInputChannel extends InputChannel implements BufferAvailabilit private final Deque<BufferAndBacklog> toBeConsumedBuffers = new ArrayDeque<>(); /** - * Buffers migrated from {@code RecoveredInputChannel}, kept separately from {@link + * Buffers delivered from {@code RecoveredInputChannel}, kept separately from {@link * #toBeConsumedBuffers} so that recovery semantics (priority event interleaving, checkpoint - * inflight persistence) do not leak into the FullyFilledBuffer split path. + * inflight persistence) do not leak into the FullyFilledBuffer split path. Holds recovered + * buffers plus the {@code RecoveryCheckpointBarrier} and {@code EndOfFetchedChannelStateEvent} + * sentinels. The deque object is its own monitor; {@link #inRecovery} and {@link + * #recoverySequenceNumber} are guarded by it too. */ - private final Deque<BufferAndBacklog> recoveredBuffers = new ArrayDeque<>(); + private final Deque<Buffer> recoveredBuffers = new ArrayDeque<>(); /** - * Flag indicating whether there is a pending priority event (e.g., checkpoint barrier) in the - * subpartitionView that should be consumed before recoveredBuffers. This is set by {@link - * #notifyPriorityEvent} and checked in {@link #getNextBuffer()}. + * Whether the channel is still replaying recovered state. Starts {@code false} for channels + * that do not need recovery and is flipped to {@code false} the moment the consume path polls + * the {@code EndOfFetchedChannelStateEvent} sentinel appended after the last recovered buffer + * (see {@link #onRecoveredStateConsumed()}). While {@code true} the consume path serves + * recovered buffers and does not poll ordinary upstream data. + */ + @GuardedBy("recoveredBuffers") + private boolean inRecovery; + + private final CompletableFuture<Void> stateConsumedFuture = new CompletableFuture<>(); + + /** + * Sequence number assigned to recovered buffers, starting at {@link Integer#MIN_VALUE}, + * consistent with {@link RecoveredInputChannel}. + */ + private int recoverySequenceNumber = Integer.MIN_VALUE; + + @Nullable private final BufferManager bufferManager; + + private final int networkBuffersPerChannel; + + private final boolean needsRecovery; + + /** + * Whether a priority event (e.g., checkpoint barrier) is pending in {@code subpartitionView} + * and must be consumed before {@code recoveredBuffers}. Volatile because it is written by the + * network thread and read by the task thread. */ private volatile boolean hasPendingPriorityEvent = false; + /** + * One-shot latch that opens once the upstream subpartition view is registered (signalled by + * {@link #requestSubpartition} or by {@link #releaseAllResources()}). Recovery-side awaiters + * block on it before handing off. + */ + private final CountDownLatch upstreamReady; + public LocalInputChannel( SingleInputGate inputGate, int channelIndex, @@ -108,7 +150,41 @@ public class LocalInputChannel extends InputChannel implements BufferAvailabilit Counter numBytesIn, Counter numBuffersIn, ChannelStateWriter stateWriter, - ArrayDeque<Buffer> initialRecoveredBuffers) { + int networkBuffersPerChannel, + boolean needsRecovery) { + this( + inputGate, + channelIndex, + partitionId, + consumedSubpartitionIndexSet, + partitionManager, + taskEventPublisher, + initialBackoff, + maxBackoff, + numBytesIn, + numBuffersIn, + stateWriter, + networkBuffersPerChannel, + needsRecovery, + new CountDownLatch(1)); + } + + @VisibleForTesting + LocalInputChannel( + SingleInputGate inputGate, + int channelIndex, + ResultPartitionID partitionId, + ResultSubpartitionIndexSet consumedSubpartitionIndexSet, + ResultPartitionManager partitionManager, + TaskEventPublisher taskEventPublisher, + int initialBackoff, + int maxBackoff, + Counter numBytesIn, + Counter numBuffersIn, + ChannelStateWriter stateWriter, + int networkBuffersPerChannel, + boolean needsRecovery, + CountDownLatch upstreamReady) { super( inputGate, @@ -122,54 +198,247 @@ public class LocalInputChannel extends InputChannel implements BufferAvailabilit this.partitionManager = checkNotNull(partitionManager); this.taskEventPublisher = checkNotNull(taskEventPublisher); - this.channelStatePersister = new ChannelStatePersister(stateWriter, getChannelInfo()); - - // Migrate recovered buffers from RecoveredInputChannel if provided. - // These buffers have been filtered but not yet consumed by the Task. - if (!initialRecoveredBuffers.isEmpty()) { - final int expectedCount = initialRecoveredBuffers.size(); - // Sequence number starts at Integer.MIN_VALUE, consistent with RecoveredInputChannel. - int seqNum = Integer.MIN_VALUE; - while (!initialRecoveredBuffers.isEmpty()) { - Buffer buffer = initialRecoveredBuffers.poll(); - // Determine next data type based on the next buffer in the queue - Buffer.DataType nextDataType = - initialRecoveredBuffers.isEmpty() - ? Buffer.DataType.NONE - : initialRecoveredBuffers.peek().getDataType(); - // buffersInBacklog is set to 0 as these are recovered buffers - BufferAndBacklog bufferAndBacklog = - new BufferAndBacklog(buffer, 0, nextDataType, seqNum++); - recoveredBuffers.add(bufferAndBacklog); - } - checkState( - recoveredBuffers.size() == expectedCount, - "Buffer migration failed: expected %s buffers but got %s", - expectedCount, - recoveredBuffers.size()); + this.channelStatePersister = + new ChannelStatePersister(checkNotNull(stateWriter), getChannelInfo()); + this.inRecovery = needsRecovery; + this.bufferManager = + needsRecovery + ? new BufferManager(inputGate.getMemorySegmentProvider(), this, 0, true) + : null; + this.networkBuffersPerChannel = networkBuffersPerChannel; + this.needsRecovery = needsRecovery; + this.upstreamReady = checkNotNull(upstreamReady); + if (!needsRecovery) { + stateConsumedFuture.complete(null); + } + } + + @Override + void setup() throws IOException { + if (needsRecovery && networkBuffersPerChannel > 0) { + bufferManager.requestExclusiveBuffers(networkBuffersPerChannel); } } // ------------------------------------------------------------------------ - // Consume + // RecoverableInputChannel implementation // ------------------------------------------------------------------------ - public void checkpointStarted(CheckpointBarrier barrier) throws CheckpointException { - // Collect inflight buffers from recoveredBuffers to be persisted. - // These are recovered buffers that have not been consumed yet when the checkpoint - // barrier arrives. - List<Buffer> inflightBuffers = new ArrayList<>(); - for (BufferAndBacklog bufferAndBacklog : recoveredBuffers) { - if (bufferAndBacklog.buffer().isBuffer()) { - inflightBuffers.add(bufferAndBacklog.buffer().retainBuffer()); + @Override + public void onRecoveredStateBuffer(Buffer buffer) { + boolean wasEmpty; + synchronized (recoveredBuffers) { + if (isReleased) { + buffer.recycleBuffer(); + return; } + // Migrate recovered buffers from RecoveredInputChannel. These buffers have been + // filtered but not yet consumed by the Task. + wasEmpty = offerRecoveredBuffer(buffer); + } + if (wasEmpty) { + notifyChannelNonEmpty(); + } + } + + @Override + public void finishRecoveredBufferDelivery() throws IOException, InterruptedException { + upstreamReady.await(); + boolean wasEmpty; + synchronized (recoveredBuffers) { + // A release may have opened the latch instead of the subpartition view; bail out so we + // never append to a queue that releaseAllResources() already cleared. + if (isReleased) { + return; + } + checkState(inRecovery, "Recovery delivery already finished."); + // Append the sentinel after the last recovered buffer. The consume path flips out of + // recovery only once it polls this sentinel, guaranteeing all recovered buffers are + // consumed first. + wasEmpty = + offerRecoveredBuffer( + EventSerializer.toBuffer( + EndOfFetchedChannelStateEvent.INSTANCE, false)); + } + if (wasEmpty) { + notifyChannelNonEmpty(); + } + } + + @Override + public Buffer requestRecoveryBufferBlocking() throws InterruptedException, IOException { + checkState( + bufferManager != null, + "requestRecoveryBufferBlocking called on a Local channel constructed with" + + " needsRecovery=false"); + upstreamReady.await(); + // If a release opened the latch instead of the subpartition view, requestBufferBlocking() + // detects the released channel and throws CancelTaskException. + return bufferManager.requestBufferBlocking(); + } + + @Override + public void insertRecoveryCheckpointBarrierIfInRecovery(long checkpointId) throws IOException { + boolean wasEmpty = false; + synchronized (recoveredBuffers) { + if (!isReleased && inRecovery) { + wasEmpty = + offerRecoveredBuffer( + EventSerializer.toBuffer( + new RecoveryCheckpointBarrier(checkpointId), false)); + } + } + if (wasEmpty) { + notifyChannelNonEmpty(); + } + } + + /** + * Flips out of recovery the moment the consume path polls the {@code + * EndOfFetchedChannelStateEvent} sentinel, i.e. once all recovered buffers have been consumed. + * Live upstream data may flow again afterwards. + */ + @Override + public void onRecoveredStateConsumed() { + synchronized (recoveredBuffers) { + checkState(inRecovery, "Recovery already finished."); + inRecovery = false; } - channelStatePersister.startPersisting(barrier.getId(), inflightBuffers); + notifyChannelNonEmpty(); + stateConsumedFuture.complete(null); } @Override public CompletableFuture<Void> getStateConsumedFuture() { - return FutureUtils.completedVoidFuture(); + return stateConsumedFuture; + } + + /** + * Appends a recovered buffer (or {@code RecoveryCheckpointBarrier} / {@code + * EndOfFetchedChannelStateEvent} sentinel) to {@link #recoveredBuffers}. + * + * @return {@code true} iff {@link #recoveredBuffers} transitioned from empty to non-empty. + */ + private boolean offerRecoveredBuffer(Buffer buffer) { + assert Thread.holdsLock(recoveredBuffers); + checkState(inRecovery, "Push into recovered buffers after recovery finished."); + boolean wasEmpty = recoveredBuffers.isEmpty(); + recoveredBuffers.add(buffer); + return wasEmpty; + } + + private int nextRecoverySequenceNumber() { + assert Thread.holdsLock(recoveredBuffers); + return recoverySequenceNumber++; + } + + /** + * Walks {@link #recoveredBuffers} up to the {@link RecoveryCheckpointBarrier} sentinel matching + * {@code checkpointId}, retaining each pre-barrier recovered data buffer and removing the + * sentinel. A barrier for an earlier (subsumed) checkpoint encountered on the way is logged and + * dropped; a barrier for a later checkpoint is an ordering violation. + * + * @throws IOException if a barrier for a later checkpoint is encountered before {@code + * checkpointId}, or if no sentinel matching {@code checkpointId} is found (the snapshot + * protocol guarantees one must be present while the channel is in recovery). + */ + private List<Buffer> collectPreRecoveryBarrier(long checkpointId) throws IOException { + assert Thread.holdsLock(recoveredBuffers); + List<Buffer> retained = new ArrayList<>(); + try { + Iterator<Buffer> it = recoveredBuffers.iterator(); + while (it.hasNext()) { + Buffer b = it.next(); + RecoveryCheckpointBarrier barrier = asRecoveryCheckpointBarrier(b); + if (barrier != null) { + long barrierId = barrier.getCheckpointId(); + if (barrierId == checkpointId) { + it.remove(); + b.recycleBuffer(); + return retained; + } + if (barrierId > checkpointId) { + throw new IOException( + "Found RecoveryCheckpointBarrier for a later checkpoint " + + barrierId + + " before the target checkpoint " + + checkpointId + + " in recoveredBuffers for channel " + + getChannelInfo()); + } + // barrierId < checkpointId: the checkpoint was subsumed; drop its stale + // barrier and keep scanning for the target. + LOG.warn( + "Discarding subsumed RecoveryCheckpointBarrier for checkpoint {} while " + + "collecting checkpoint {} on channel {}.", + barrierId, + checkpointId, + getChannelInfo()); + it.remove(); + b.recycleBuffer(); + continue; + } + if (b.isBuffer()) { + retained.add(b.retainBuffer()); + } + } + } catch (IOException e) { + releaseRetainedBuffers(retained); + throw e; + } + releaseRetainedBuffers(retained); + throw new IOException( + "Missing RecoveryCheckpointBarrier for checkpoint " + + checkpointId + + " in recoveredBuffers for channel " + + getChannelInfo()); + } + + private static void releaseRetainedBuffers(List<Buffer> retained) { + for (Buffer buffer : retained) { + buffer.recycleBuffer(); + } + } + + @Nullable + private static RecoveryCheckpointBarrier asRecoveryCheckpointBarrier(Buffer b) + throws IOException { + if (b.isBuffer()) { + return null; + } + AbstractEvent event = + EventSerializer.fromBuffer(b, RecoveryCheckpointBarrier.class.getClassLoader()); + b.setReaderIndex(0); + return event instanceof RecoveryCheckpointBarrier + ? (RecoveryCheckpointBarrier) event + : null; + } + + // ------------------------------------------------------------------------ + // Consume + // ------------------------------------------------------------------------ + + @Override + public void checkpointStarted(CheckpointBarrier barrier) throws CheckpointException { + try { + List<Buffer> toPersist; + synchronized (recoveredBuffers) { + if (inRecovery) { + // Collect inflight buffers from recoveredBuffers to be persisted. These are + // recovered buffers that have not been consumed yet when the checkpoint barrier + // arrives. + toPersist = collectPreRecoveryBarrier(barrier.getId()); + } else { + toPersist = Collections.emptyList(); + } + } + channelStatePersister.startPersisting(barrier.getId(), toPersist); + } catch (IOException e) { + throw new CheckpointException( + "Failed to extract recovered buffers for checkpoint " + barrier.getId(), + CheckpointFailureReason.CHECKPOINT_DECLINED, + e); + } } public void checkpointStopped(long checkpointId) { @@ -213,6 +482,7 @@ public class LocalInputChannel extends InputChannel implements BufferAvailabilit this.subpartitionView = null; } else { notifyDataAvailable = true; + upstreamReady.countDown(); } } catch (PartitionNotFoundException notFound) { if (increaseBackoff()) { @@ -289,12 +559,35 @@ public class LocalInputChannel extends InputChannel implements BufferAvailabilit public Optional<BufferAndAvailability> getNextBuffer() throws IOException { checkError(); - if (!recoveredBuffers.isEmpty()) { - return getNextRecoveredBuffer(); + // Read inRecovery and poll the recovered buffer under a single lock acquisition to avoid + // grabbing the monitor twice on the hot path. + boolean inRecovery; + Buffer recoveredBuf = null; + synchronized (recoveredBuffers) { + inRecovery = this.inRecovery; + if (inRecovery && !hasPendingPriorityEvent && !recoveredBuffers.isEmpty()) { + recoveredBuf = recoveredBuffers.poll(); + } + } + + if (inRecovery) { + // Always return an already-polled recovered buffer first: hasPendingPriorityEvent may + // be flipped to true by a concurrent notifyPriorityEvent() after the poll, and + // re-reading + // it here would otherwise drop this buffer. A pending priority event is served on the + // next getNextBuffer() call instead. + if (recoveredBuf != null) { + return wrapRecoveredBufferAsAvailability(recoveredBuf); + } + if (hasPendingPriorityEvent) { + return pullPriorityFromSubpartitionView(); + } + // Drain not finished yet; block normal upstream data until delivery completes. + return Optional.empty(); } if (!toBeConsumedBuffers.isEmpty()) { - return Optional.of(getBufferAndAvailability(toBeConsumedBuffers.removeFirst())); + return getBufferAndAvailability(toBeConsumedBuffers.removeFirst()); } ResultSubpartitionView subpartitionView = this.subpartitionView; @@ -350,79 +643,102 @@ public class LocalInputChannel extends InputChannel implements BufferAvailabilit seq++)); } - return Optional.of(getBufferAndAvailability(toBeConsumedBuffers.removeFirst())); + return getBufferAndAvailability(toBeConsumedBuffers.removeFirst()); } - BufferAndAvailability bufferAndAvailability = getBufferAndAvailability(next); - channelStatePersister.checkForBarrier(bufferAndAvailability.buffer()); - channelStatePersister.maybePersist(bufferAndAvailability.buffer()); - return Optional.of(bufferAndAvailability); + return getBufferAndAvailability(next); } - /** - * Consumes the next buffer from recoveredBuffers, handling pending priority events and dynamic - * availability detection for the last recovered buffer. - */ - private Optional<BufferAndAvailability> getNextRecoveredBuffer() throws IOException { - // If there is a pending priority event (e.g., unaligned checkpoint barrier), fetch it - // from subpartitionView first, skipping recoveredBuffers. This ensures priority - // events are processed immediately even when there are pending recovered buffers. - if (hasPendingPriorityEvent) { - checkState(subpartitionView != null, "No subpartition view available"); - BufferAndBacklog next = subpartitionView.getNextBuffer(); - checkState( - next != null && next.buffer().getDataType().hasPriority(), - "Expected priority event, but got %s", - next == null ? "null" : next.buffer().getDataType()); - - // Check for barrier to update channel state persister. - // Note: maybePersist is not needed for barriers as they are not regular data buffers. - channelStatePersister.checkForBarrier(next.buffer()); - - Buffer.DataType expectedNextDataType = next.getNextDataType(); - if (!expectedNextDataType.hasPriority()) { - // Reset hasPendingPriorityEvent to false if no more priority event - hasPendingPriorityEvent = false; - if (!recoveredBuffers.isEmpty()) { - // Correct nextDataType: if recoveredBuffers is not empty, the actual next - // element to consume is from recoveredBuffers, not from subpartitionView - expectedNextDataType = recoveredBuffers.peek().buffer().getDataType(); - } - } + private Optional<BufferAndAvailability> pullPriorityFromSubpartitionView() throws IOException { + // If there is a pending priority event (e.g., unaligned checkpoint barrier), fetch it from + // subpartitionView first, skipping recoveredBuffers. This ensures priority events are + // processed immediately even when there are pending recovered buffers. + checkState(subpartitionView != null, "No subpartition view available"); + BufferAndBacklog next = subpartitionView.getNextBuffer(); + checkState( + next != null && next.buffer().getDataType().hasPriority(), + "Expected priority event, but got %s", + next == null ? "null" : next.buffer().getDataType()); + + // Check for barrier to update channel state persister. Note: maybePersist is not needed for + // barriers as they are not regular data buffers. + channelStatePersister.checkForBarrier(next.buffer()); + + Buffer.DataType expectedNextDataType = next.getNextDataType(); + if (!expectedNextDataType.hasPriority()) { + // Reset hasPendingPriorityEvent to false if no more priority event. + hasPendingPriorityEvent = false; + // Correct nextDataType: if recoveredBuffers is not empty, the actual next element to + // consume is from recoveredBuffers, not from subpartitionView. + expectedNextDataType = peekNextDataType(next.getNextDataType()); + } - return Optional.of( - getBufferAndAvailability( - new BufferAndBacklog( - next.buffer(), - next.buffersInBacklog(), - expectedNextDataType, - next.getSequenceNumber()))); - } - - BufferAndBacklog next = recoveredBuffers.removeFirst(); - - // If this is the last recovered buffer and nextDataType is NONE, - // dynamically check if subpartitionView has data available. - // The last buffer's nextDataType was preset to NONE during construction, - // but subpartitionView may already have data available. - if (recoveredBuffers.isEmpty() - && next.getNextDataType() == Buffer.DataType.NONE - && subpartitionView != null) { - ResultSubpartitionView.AvailabilityWithBacklog availability = - subpartitionView.getAvailabilityAndBacklog(true); - if (availability.isAvailable()) { - next = - new BufferAndBacklog( - next.buffer(), - availability.getBacklog(), - Buffer.DataType.DATA_BUFFER, - next.getSequenceNumber()); + return Optional.of( + new BufferAndAvailability( + next.buffer(), + expectedNextDataType, + next.buffersInBacklog(), + next.getSequenceNumber())); + } + + private Optional<BufferAndAvailability> wrapRecoveredBufferAsAvailability(Buffer buf) + throws IOException { + if (buf instanceof FileRegionBuffer) { + buf = ((FileRegionBuffer) buf).readInto(inputGate.getUnpooledSegment()); + } + if (buf instanceof CompositeBuffer) { + buf = ((CompositeBuffer) buf).getFullBufferData(inputGate.getUnpooledSegment()); + } + + numBytesIn.inc(buf.readableBytes()); + numBuffersIn.inc(); + + ResultSubpartitionView view = subpartitionView; + Buffer.DataType upstreamProbe; + if (view != null && view.getAvailabilityAndBacklog(true).isAvailable()) { + upstreamProbe = Buffer.DataType.DATA_BUFFER; + } else { + upstreamProbe = Buffer.DataType.NONE; + } + + int sequenceNumber; + synchronized (recoveredBuffers) { + Buffer.DataType nextDataType = peekNextDataType(upstreamProbe); + sequenceNumber = nextRecoverySequenceNumber(); + NetworkActionsLogger.traceInput( + "LocalInputChannel#getNextBuffer", + buf, + inputGate.getOwningTaskName(), + channelInfo, + channelStatePersister, + sequenceNumber); + // buffersInBacklog is set to 0 as these are recovered buffers. + return Optional.of(new BufferAndAvailability(buf, nextDataType, 0, sequenceNumber)); + } + } + + private Buffer.DataType peekNextDataType(Buffer.DataType nextDataTypeOnUpstream) { + synchronized (recoveredBuffers) { + if (!recoveredBuffers.isEmpty()) { + return recoveredBuffers.peek().getDataType(); + } + if (inRecovery) { + // If this is the last currently available recovered buffer, hide upstream data + // until the EndOfFetchedChannelStateEvent sentinel flips the channel out of + // recovery. The last buffer's nextDataType is effectively NONE while the drain can + // still append more recovered buffers. + return Buffer.DataType.NONE; } } - return Optional.of(getBufferAndAvailability(next)); + // If this is the last recovered buffer after delivery finished, dynamically check if + // subpartitionView has data available. The last buffer's nextDataType may have been NONE + // while recovered data was still being delivered, but subpartitionView may already have + // data + // available now. + return nextDataTypeOnUpstream; } - private BufferAndAvailability getBufferAndAvailability(BufferAndBacklog next) + private Optional<BufferAndAvailability> getBufferAndAvailability(BufferAndBacklog next) throws IOException { Buffer buffer = next.buffer(); if (buffer instanceof FileRegionBuffer) { @@ -435,6 +751,8 @@ public class LocalInputChannel extends InputChannel implements BufferAvailabilit numBytesIn.inc(buffer.readableBytes()); numBuffersIn.inc(); + channelStatePersister.checkForBarrier(buffer); + channelStatePersister.maybePersist(buffer); NetworkActionsLogger.traceInput( "LocalInputChannel#getNextBuffer", buffer, @@ -442,8 +760,12 @@ public class LocalInputChannel extends InputChannel implements BufferAvailabilit channelInfo, channelStatePersister, next.getSequenceNumber()); - return new BufferAndAvailability( - buffer, next.getNextDataType(), next.buffersInBacklog(), next.getSequenceNumber()); + return Optional.of( + new BufferAndAvailability( + buffer, + next.getNextDataType(), + next.buffersInBacklog(), + next.getSequenceNumber())); } @Override @@ -454,7 +776,7 @@ public class LocalInputChannel extends InputChannel implements BufferAvailabilit @Override public void notifyPriorityEvent(int prioritySequenceNumber) { // Set flag so that getNextBuffer() knows to fetch priority event from subpartitionView - // before consuming toBeConsumedBuffers. + // before consuming recoveredBuffers. hasPendingPriorityEvent = true; super.notifyPriorityEvent(prioritySequenceNumber); } @@ -525,23 +847,35 @@ public class LocalInputChannel extends InputChannel implements BufferAvailabilit if (!isReleased) { isReleased = true; + // Unblock any thread awaiting upstreamReady (drain still in flight) so it falls + // through and observes the released state instead of deadlocking. + upstreamReady.countDown(); + + // Recovery will never be consumed on a released channel; unblock anyone gating on it. + stateConsumedFuture.completeExceptionally(new CancelTaskException("Channel released.")); + ResultSubpartitionView view = subpartitionView; if (view != null) { view.releaseAllResources(); subpartitionView = null; } - // Release any remaining buffers in recoveredBuffers (migrated recovered buffers - // not yet consumed) and toBeConsumedBuffers (FullyFilledBuffer partial splits) - // to avoid memory leak. - for (BufferAndBacklog bufferAndBacklog : recoveredBuffers) { - bufferAndBacklog.buffer().recycleBuffer(); + // Release any remaining buffers in recoveredBuffers (migrated recovered buffers not yet + // consumed) and toBeConsumedBuffers (FullyFilledBuffer partial splits) to avoid memory + // leak. + synchronized (recoveredBuffers) { + for (Buffer buffer : recoveredBuffers) { + buffer.recycleBuffer(); + } + recoveredBuffers.clear(); } - recoveredBuffers.clear(); for (BufferAndBacklog bufferAndBacklog : toBeConsumedBuffers) { bufferAndBacklog.buffer().recycleBuffer(); } toBeConsumedBuffers.clear(); + if (bufferManager != null) { + bufferManager.releaseAllBuffers(new ArrayDeque<>()); + } } } diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/LocalRecoveredInputChannel.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/LocalRecoveredInputChannel.java index bdde2244f38..8ec00582215 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/LocalRecoveredInputChannel.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/LocalRecoveredInputChannel.java @@ -65,6 +65,8 @@ public class LocalRecoveredInputChannel extends RecoveredInputChannel { @Override protected InputChannel toInputChannelInternal(ArrayDeque<Buffer> remainingBuffers) { + // remainingBuffers is unused: with the flag off the queue is always empty at conversion, + // and constructor migration is retired by the conversion rework later in this PR. return new LocalInputChannel( inputGate, getChannelIndex(), @@ -77,6 +79,7 @@ public class LocalRecoveredInputChannel extends RecoveredInputChannel { numBytesIn, numBuffersIn, channelStateWriter, - remainingBuffers); + networkBuffersPerChannel, + false); } } diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/UnknownInputChannel.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/UnknownInputChannel.java index ad59dd38db8..78d6f024fd4 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/UnknownInputChannel.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/UnknownInputChannel.java @@ -208,7 +208,8 @@ class UnknownInputChannel extends InputChannel implements ChannelStateHolder { metrics.getNumBytesInLocalCounter(), metrics.getNumBuffersInLocalCounter(), channelStateWriter == null ? ChannelStateWriter.NO_OP : channelStateWriter, - new ArrayDeque<>()); + networkBuffersPerChannel, + false); } @Override diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/InputChannelBuilder.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/InputChannelBuilder.java index 08f65d9fe72..1130c13bd4f 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/InputChannelBuilder.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/InputChannelBuilder.java @@ -34,7 +34,7 @@ import org.apache.flink.runtime.io.network.partition.ResultPartitionManager; import org.apache.flink.runtime.io.network.partition.ResultSubpartitionIndexSet; import java.net.InetSocketAddress; -import java.util.ArrayDeque; +import java.util.concurrent.CountDownLatch; import static org.apache.flink.runtime.io.network.partition.consumer.SingleInputGateTest.TestingResultPartitionManager; @@ -56,6 +56,8 @@ public class InputChannelBuilder { private int maxBackoff = 0; private int partitionRequestListenerTimeout = 0; private int networkBuffersPerChannel = 2; + private boolean needsRecovery = false; + private CountDownLatch upstreamReady = new CountDownLatch(1); private InputChannelMetrics metrics = InputChannelTestUtils.newUnregisteredInputChannelMetrics(); @@ -115,6 +117,16 @@ public class InputChannelBuilder { return this; } + public InputChannelBuilder setNeedsRecovery(boolean needsRecovery) { + this.needsRecovery = needsRecovery; + return this; + } + + public InputChannelBuilder setUpstreamReady(CountDownLatch upstreamReady) { + this.upstreamReady = upstreamReady; + return this; + } + public InputChannelBuilder setMetrics(InputChannelMetrics metrics) { this.metrics = metrics; return this; @@ -166,7 +178,9 @@ public class InputChannelBuilder { metrics.getNumBytesInLocalCounter(), metrics.getNumBuffersInLocalCounter(), stateWriter, - new ArrayDeque<>()); + networkBuffersPerChannel, + needsRecovery, + upstreamReady); } public RemoteInputChannel buildRemoteChannel(SingleInputGate inputGate) { diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/LocalInputChannelTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/LocalInputChannelTest.java index 44c4b48cb7a..d7dfa095818 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/LocalInputChannelTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/LocalInputChannelTest.java @@ -19,10 +19,13 @@ package org.apache.flink.runtime.io.network.partition.consumer; import org.apache.flink.metrics.SimpleCounter; +import org.apache.flink.runtime.checkpoint.CheckpointException; +import org.apache.flink.runtime.checkpoint.CheckpointFailureReason; import org.apache.flink.runtime.checkpoint.CheckpointOptions; import org.apache.flink.runtime.checkpoint.CheckpointType; import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter; import org.apache.flink.runtime.checkpoint.channel.RecordingChannelStateWriter; +import org.apache.flink.runtime.checkpoint.channel.RecoveryCheckpointBarrier; import org.apache.flink.runtime.execution.CancelTaskException; import org.apache.flink.runtime.io.disk.NoOpFileChannelManager; import org.apache.flink.runtime.io.network.TaskEventDispatcher; @@ -61,7 +64,6 @@ import org.junit.jupiter.api.Test; import org.mockito.stubbing.Answer; import java.io.IOException; -import java.util.ArrayDeque; import java.util.Collections; import java.util.List; import java.util.Optional; @@ -69,6 +71,7 @@ import java.util.Timer; import java.util.TimerTask; import java.util.concurrent.Callable; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicInteger; @@ -669,7 +672,7 @@ class LocalInputChannelTest { @Test void testGetBuffersInUseCountIncludesToBeConsumedBuffers() throws Exception { - // given: Local input channel with recovered buffers in toBeConsumedBuffers + // given: Local input channel with recovered buffers in the recovery queue ResultSubpartitionView subpartitionView = InputChannelTestUtils.createResultSubpartitionView( createFilledFinishedBufferConsumer(4096), @@ -678,12 +681,6 @@ class LocalInputChannelTest { new TestingResultPartitionManager(subpartitionView); final SingleInputGate inputGate = createSingleInputGate(1); - // Create 3 recovered buffers - ArrayDeque<Buffer> recoveredBuffers = new ArrayDeque<>(); - recoveredBuffers.add(TestBufferFactory.createBuffer(32)); - recoveredBuffers.add(TestBufferFactory.createBuffer(32)); - recoveredBuffers.add(TestBufferFactory.createBuffer(32)); - final LocalInputChannel localChannel = new LocalInputChannel( inputGate, @@ -697,10 +694,16 @@ class LocalInputChannelTest { new SimpleCounter(), new SimpleCounter(), ChannelStateWriter.NO_OP, - recoveredBuffers); + 2, + true); inputGate.setInputChannels(localChannel); + // Create 3 recovered buffers + localChannel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(32)); + localChannel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(32)); + localChannel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(32)); + // then: Before requesting subpartitions, buffers in use should include recovered buffers assertThat(localChannel.getBuffersInUseCount()).isEqualTo(3); assertThat(localChannel.unsynchronizedGetNumberOfQueuedBuffers()).isEqualTo(3); @@ -718,10 +721,7 @@ class LocalInputChannelTest { // given: LocalInputChannel with recovered buffers migrated from RecoveredInputChannel SingleInputGate inputGate = createSingleInputGate(1); - ArrayDeque<Buffer> recoveredBuffers = new ArrayDeque<>(); - recoveredBuffers.add(TestBufferFactory.createBuffer(10)); - recoveredBuffers.add(TestBufferFactory.createBuffer(20)); - + CountDownLatch upstreamReady = new CountDownLatch(1); LocalInputChannel channel = new LocalInputChannel( inputGate, @@ -735,10 +735,17 @@ class LocalInputChannelTest { new SimpleCounter(), new SimpleCounter(), ChannelStateWriter.NO_OP, - recoveredBuffers); + 2, + true, + upstreamReady); inputGate.setInputChannels(channel); + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(10)); + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(20)); + upstreamReady.countDown(); + channel.finishRecoveredBufferDelivery(); + // then: Can read recovered buffers even before requestSubpartitions() Optional<InputChannel.BufferAndAvailability> first = channel.getNextBuffer(); assertThat(first).isPresent(); @@ -755,11 +762,6 @@ class LocalInputChannelTest { // given: Local input channel with recovered buffers SingleInputGate inputGate = new SingleInputGateBuilder().build(); - ArrayDeque<Buffer> recoveredBuffers = new ArrayDeque<>(); - recoveredBuffers.add(TestBufferFactory.createBuffer(10)); - recoveredBuffers.add(TestBufferFactory.createBuffer(20)); - recoveredBuffers.add(TestBufferFactory.createBuffer(30)); - RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); LocalInputChannel channel = @@ -775,53 +777,51 @@ class LocalInputChannelTest { new SimpleCounter(), new SimpleCounter(), stateWriter, - recoveredBuffers); + 2, + true); inputGate.setInputChannels(channel); - // when: Checkpoint is started + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(10)); + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(20)); + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(30)); + // Mirror snapshotAndInsertBarriers: push the sentinel before + // checkpointStarted scans recoveredQueue. + channel.onRecoveredStateBuffer( + EventSerializer.toBuffer(new RecoveryCheckpointBarrier(1L), false)); + CheckpointOptions options = CheckpointOptions.unaligned(CheckpointType.CHECKPOINT, getDefault()); stateWriter.start(1L, options); CheckpointBarrier barrier = new CheckpointBarrier(1L, 0L, options); + // when: Checkpoint is started channel.checkpointStarted(barrier); - // then: All 3 recovered buffers should be persisted as inflight data List<Buffer> persistedBuffers = stateWriter.getAddedInput().get(channel.getChannelInfo()); + // then: All 3 recovered buffers should be persisted as inflight data assertThat(persistedBuffers).isNotNull().hasSize(3); assertThat(persistedBuffers.stream().mapToInt(Buffer::getSize).toArray()) .containsExactly(10, 20, 30); } + /** Port of the FLINK-40016 double-persist regression test to the push-based recovery API. */ @Test void testRecoveredBuffersNotPersistedAgainWhenConsumedDuringCheckpoint() throws Exception { - // given: Channel with 2 recovered buffers + // given: Channel with 2 recovered buffers, followed by the RecoveryCheckpointBarrier + // sentinel that snapshotAndInsertBarriers pushes before checkpointStarted scans + // recoveredBuffers SingleInputGate inputGate = new SingleInputGateBuilder().build(); - - ArrayDeque<Buffer> recoveredBuffers = new ArrayDeque<>(); - recoveredBuffers.add(TestBufferFactory.createBuffer(10)); - recoveredBuffers.add(TestBufferFactory.createBuffer(20)); - RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); - - LocalInputChannel channel = - new LocalInputChannel( - inputGate, - 0, - new ResultPartitionID(), - new ResultSubpartitionIndexSet(0), - new ResultPartitionManager(), - new TaskEventDispatcher(), - 0, - 0, - new SimpleCounter(), - new SimpleCounter(), - stateWriter, - recoveredBuffers); - + LocalInputChannel channel = newPushOnlyLocalChannel(inputGate, stateWriter); inputGate.setInputChannels(channel); - // when: Checkpoint starts — checkpointStarted spills recovered buffers as inflight data + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(10)); + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(20)); + channel.onRecoveredStateBuffer( + EventSerializer.toBuffer(new RecoveryCheckpointBarrier(1L), false)); + + // when: Checkpoint starts — checkpointStarted persists the pre-barrier recovered buffers + // as inflight data CheckpointOptions options = CheckpointOptions.unaligned(CheckpointType.CHECKPOINT, getDefault()); stateWriter.start(1L, options); @@ -830,15 +830,22 @@ class LocalInputChannelTest { // then: exactly 2 buffers persisted so far assertThat(stateWriter.getAddedInput().get(channel.getChannelInfo())).hasSize(2); - // when: Consume the recovered buffers while checkpoint is still in BARRIER_PENDING state - channel.getNextBuffer(); - channel.getNextBuffer(); + // when: Consume the recovered buffers while the checkpoint is still in BARRIER_PENDING + // state + Optional<InputChannel.BufferAndAvailability> first = channel.getNextBuffer(); + assertThat(first).isPresent(); + assertThat(first.get().buffer().getSize()).isEqualTo(10); + Optional<InputChannel.BufferAndAvailability> second = channel.getNextBuffer(); + assertThat(second).isPresent(); + assertThat(second.get().buffer().getSize()).isEqualTo(20); // then: total persisted count must still be 2 — recovered buffers must not be // re-persisted via maybePersist() when consumed (that would make it 4 without the fix) - assertThat(stateWriter.getAddedInput().get(channel.getChannelInfo())) + List<Buffer> persisted = stateWriter.getAddedInput().get(channel.getChannelInfo()); + assertThat(persisted) .as("recovered buffers must not be persisted again when consumed") .hasSize(2); + assertThat(persisted.stream().mapToInt(Buffer::getSize).toArray()).containsExactly(10, 20); } @Test @@ -872,9 +879,6 @@ class LocalInputChannelTest { // given: Local input channel with recovered buffers but NO subpartition view initialized SingleInputGate inputGate = new SingleInputGateBuilder().build(); - ArrayDeque<Buffer> recoveredBuffers = new ArrayDeque<>(); - recoveredBuffers.add(TestBufferFactory.createBuffer(10)); - LocalInputChannel channel = new LocalInputChannel( inputGate, @@ -888,9 +892,11 @@ class LocalInputChannelTest { new SimpleCounter(), new SimpleCounter(), ChannelStateWriter.NO_OP, - recoveredBuffers); + 2, + true); inputGate.setInputChannels(channel); + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(10)); // Do NOT call channel.requestSubpartitions() — subpartitionView stays null channel.notifyPriorityEvent(0); @@ -1010,11 +1016,6 @@ class LocalInputChannelTest { TestingResultPartitionManager partitionManager = new TestingResultPartitionManager(subpartitionView); - ArrayDeque<Buffer> recoveredBuffers = new ArrayDeque<>(); - for (int size : recoveredBufferSizes) { - recoveredBuffers.add(TestBufferFactory.createBuffer(size)); - } - LocalInputChannel channel = new LocalInputChannel( inputGate, @@ -1028,9 +1029,15 @@ class LocalInputChannelTest { new SimpleCounter(), new SimpleCounter(), stateWriter, - recoveredBuffers); + 2, + true); inputGate.setInputChannels(channel); + + for (int size : recoveredBufferSizes) { + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(size)); + } + channel.requestSubpartitions(); return new ChannelAndSubpartition(channel, subpartition); @@ -1046,6 +1053,351 @@ class LocalInputChannelTest { } } + // --------------------------------------------------------------------------------------------- + // RecoverableInputChannel push-based recovery tests + // --------------------------------------------------------------------------------------------- + + private static LocalInputChannel newPushOnlyLocalChannel( + SingleInputGate inputGate, ChannelStateWriter stateWriter) { + return newPushOnlyLocalChannel(inputGate, stateWriter, new CountDownLatch(1)); + } + + private static LocalInputChannel newPushOnlyLocalChannel( + SingleInputGate inputGate, + ChannelStateWriter stateWriter, + CountDownLatch upstreamReady) { + return new LocalInputChannel( + inputGate, + 0, + new ResultPartitionID(), + new ResultSubpartitionIndexSet(0), + new ResultPartitionManager(), + new TaskEventDispatcher(), + 0, + 0, + new SimpleCounter(), + new SimpleCounter(), + stateWriter, + 2, + true, + upstreamReady); + } + + @Test + void testOnRecoveredStateBufferEnqueues() throws Exception { + SingleInputGate inputGate = new SingleInputGateBuilder().build(); + LocalInputChannel channel = newPushOnlyLocalChannel(inputGate, ChannelStateWriter.NO_OP); + inputGate.setInputChannels(channel); + + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(11)); + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(22)); + + Optional<InputChannel.BufferAndAvailability> first = channel.getNextBuffer(); + assertThat(first).isPresent(); + assertThat(first.get().buffer().getSize()).isEqualTo(11); + Optional<InputChannel.BufferAndAvailability> second = channel.getNextBuffer(); + assertThat(second).isPresent(); + assertThat(second.get().buffer().getSize()).isEqualTo(22); + } + + @Test + void testOnRecoveredStateBufferOnReleasedChannelIsSilentlyRecycled() throws Exception { + SingleInputGate inputGate = new SingleInputGateBuilder().build(); + LocalInputChannel channel = newPushOnlyLocalChannel(inputGate, ChannelStateWriter.NO_OP); + inputGate.setInputChannels(channel); + channel.releaseAllResources(); + + Buffer b = TestBufferFactory.createBuffer(33); + channel.onRecoveredStateBuffer(b); + assertThat(b.isRecycled()).isTrue(); + } + + @Test + void testOnRecoveredStateBufferNotifiesChannelNonEmptyOnEmptyToNonEmptyTransition() + throws Exception { + SingleInputGate inputGate = new SingleInputGateBuilder().build(); + LocalInputChannel channel = newPushOnlyLocalChannel(inputGate, ChannelStateWriter.NO_OP); + inputGate.setInputChannels(channel); + + CompletableFuture<?> availability = inputGate.getAvailableFuture(); + assertThat(availability).isNotDone(); + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(1)); + assertThat(availability).isDone(); + } + + @Test + void testInRecoveryBoundaryFlagFalseQueueEmptyReturnsEmpty() throws Exception { + // Drive into the (flag=false, queue=empty) boundary by pushing one buffer, polling it, + // and verifying the channel does not yet expose a master-path result. + SingleInputGate inputGate = new SingleInputGateBuilder().build(); + LocalInputChannel channel = newPushOnlyLocalChannel(inputGate, ChannelStateWriter.NO_OP); + inputGate.setInputChannels(channel); + + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(1)); + channel.getNextBuffer(); + // Queue is empty; no finishRecoveredBufferDelivery was called. Without the subpartitionView + // active, the channel returns empty. + Optional<InputChannel.BufferAndAvailability> result = channel.getNextBuffer(); + assertThat(result).isNotPresent(); + } + + @Test + void testInRecoveryBoundaryFlagFalseQueueNonEmptyPolls() throws Exception { + SingleInputGate inputGate = new SingleInputGateBuilder().build(); + LocalInputChannel channel = newPushOnlyLocalChannel(inputGate, ChannelStateWriter.NO_OP); + inputGate.setInputChannels(channel); + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(7)); + Optional<InputChannel.BufferAndAvailability> r = channel.getNextBuffer(); + assertThat(r).isPresent(); + assertThat(r.get().buffer().getSize()).isEqualTo(7); + } + + @Test + void testInRecoveryBoundaryFlagTrueQueueNonEmptyPolls() throws Exception { + SingleInputGate inputGate = new SingleInputGateBuilder().build(); + CountDownLatch upstreamReady = new CountDownLatch(1); + LocalInputChannel channel = + newPushOnlyLocalChannel(inputGate, ChannelStateWriter.NO_OP, upstreamReady); + inputGate.setInputChannels(channel); + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(8)); + upstreamReady.countDown(); + channel.finishRecoveredBufferDelivery(); + Optional<InputChannel.BufferAndAvailability> r = channel.getNextBuffer(); + assertThat(r).isPresent(); + assertThat(r.get().buffer().getSize()).isEqualTo(8); + } + + @Test + void testFinishWithNoRecoveredBuffersEmitsSentinelThenFallsToMasterPath() throws Exception { + // With no recovered buffers, finish still appends the EndOfFetchedChannelStateEvent + // sentinel; getNextBuffer returns it. Consuming the sentinel flips the channel out of + // recovery, after which the master path takes over: without a subpartition view it raises + // IllegalStateException via checkAndWaitForSubpartitionView, proving the recovery branch is + // no longer swallowing the call. + SingleInputGate inputGate = new SingleInputGateBuilder().build(); + CountDownLatch upstreamReady = new CountDownLatch(1); + LocalInputChannel channel = + newPushOnlyLocalChannel(inputGate, ChannelStateWriter.NO_OP, upstreamReady); + inputGate.setInputChannels(channel); + upstreamReady.countDown(); + channel.finishRecoveredBufferDelivery(); + + Optional<InputChannel.BufferAndAvailability> sentinel = channel.getNextBuffer(); + assertThat(sentinel).isPresent(); + assertThat(EventSerializer.fromBuffer(sentinel.get().buffer(), getClass().getClassLoader())) + .isInstanceOf(EndOfFetchedChannelStateEvent.class); + + channel.onRecoveredStateConsumed(); + assertThatThrownBy(channel::getNextBuffer) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("Queried for a buffer before requesting the subpartition"); + } + + @Test + void testPriorityEventDuringRecoveryFetchedFromSubpartitionView() throws Exception { + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + ChannelAndSubpartition ctx = createChannelWithRecoveredBuffers(stateWriter, 10, 20); + + CheckpointOptions options = + CheckpointOptions.unaligned(CheckpointType.CHECKPOINT, getDefault()); + ctx.subpartition.add( + EventSerializer.toBufferConsumer(new CheckpointBarrier(1L, 0L, options), true)); + ctx.channel.notifyPriorityEvent(0); + + Optional<InputChannel.BufferAndAvailability> first = ctx.channel.getNextBuffer(); + assertThat(first).isPresent(); + assertThat(first.get().buffer().getDataType().hasPriority()).isTrue(); + } + + @Test + void testPriorityEventDuringRecoveryResetAfterNonPriority() throws Exception { + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + ChannelAndSubpartition ctx = createChannelWithRecoveredBuffers(stateWriter, 10); + + CheckpointOptions options = + CheckpointOptions.unaligned(CheckpointType.CHECKPOINT, getDefault()); + ctx.subpartition.add( + EventSerializer.toBufferConsumer(new CheckpointBarrier(1L, 0L, options), true)); + ctx.subpartition.add(createFilledFinishedBufferConsumer(32)); + ctx.channel.notifyPriorityEvent(0); + + Optional<InputChannel.BufferAndAvailability> priority = ctx.channel.getNextBuffer(); + assertThat(priority).isPresent(); + Optional<InputChannel.BufferAndAvailability> recovered = ctx.channel.getNextBuffer(); + assertThat(recovered).isPresent(); + assertThat(recovered.get().buffer().isBuffer()).isTrue(); + assertThat(recovered.get().buffer().getSize()).isEqualTo(10); + } + + @Test + void testCheckpointStartedScansRecoveredBuffersUpToBarrier() throws Exception { + SingleInputGate inputGate = new SingleInputGateBuilder().build(); + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + LocalInputChannel channel = newPushOnlyLocalChannel(inputGate, stateWriter); + inputGate.setInputChannels(channel); + + Buffer b1 = TestBufferFactory.createBuffer(1); + Buffer b2 = TestBufferFactory.createBuffer(2); + Buffer b3 = TestBufferFactory.createBuffer(3); + channel.onRecoveredStateBuffer(b1); + channel.onRecoveredStateBuffer(b2); + channel.onRecoveredStateBuffer( + EventSerializer.toBuffer(new RecoveryCheckpointBarrier(1L), false)); + channel.onRecoveredStateBuffer(b3); + + CheckpointOptions options = + CheckpointOptions.unaligned(CheckpointType.CHECKPOINT, getDefault()); + stateWriter.start(1L, options); + channel.checkpointStarted(new CheckpointBarrier(1L, 0L, options)); + + List<Buffer> persisted = stateWriter.getAddedInput().get(channel.getChannelInfo()); + assertThat(persisted).hasSize(2); + assertThat(persisted.stream().mapToInt(Buffer::getSize).toArray()).containsExactly(1, 2); + } + + @Test + void testCheckpointStartedDeclinesWhenRecoveryBarrierIsMissing() throws Exception { + SingleInputGate inputGate = new SingleInputGateBuilder().build(); + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + LocalInputChannel channel = newPushOnlyLocalChannel(inputGate, stateWriter); + inputGate.setInputChannels(channel); + + Buffer b1 = TestBufferFactory.createBuffer(1); + channel.onRecoveredStateBuffer(b1); + int refCntBefore = b1.refCnt(); + + CheckpointOptions options = + CheckpointOptions.unaligned(CheckpointType.CHECKPOINT, getDefault()); + stateWriter.start(1L, options); + + // The snapshot protocol guarantees a RecoveryCheckpointBarrier sentinel is present while + // the channel is in recovery, so a missing sentinel is a protocol violation: the checkpoint + // is declined and the recovered buffer is neither dropped nor persisted. + assertThatThrownBy(() -> channel.checkpointStarted(new CheckpointBarrier(1L, 0L, options))) + .isInstanceOfSatisfying( + CheckpointException.class, + e -> + assertThat(e.getCheckpointFailureReason()) + .isEqualTo(CheckpointFailureReason.CHECKPOINT_DECLINED)) + .cause() + .hasMessageContaining("Missing RecoveryCheckpointBarrier"); + assertThat(b1.refCnt()).isEqualTo(refCntBefore); + assertThat(stateWriter.getAddedInput().get(channel.getChannelInfo())).isEmpty(); + } + + @Test + void testCheckpointStartedRetainsPreBarrierBuffers() throws Exception { + SingleInputGate inputGate = new SingleInputGateBuilder().build(); + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + LocalInputChannel channel = newPushOnlyLocalChannel(inputGate, stateWriter); + inputGate.setInputChannels(channel); + + Buffer b1 = TestBufferFactory.createBuffer(1); + channel.onRecoveredStateBuffer(b1); + channel.onRecoveredStateBuffer( + EventSerializer.toBuffer(new RecoveryCheckpointBarrier(1L), false)); + + int before = b1.refCnt(); + CheckpointOptions options = + CheckpointOptions.unaligned(CheckpointType.CHECKPOINT, getDefault()); + stateWriter.start(1L, options); + channel.checkpointStarted(new CheckpointBarrier(1L, 0L, options)); + // Pre-barrier buffers are retained for the writer; their ref-count stays at or above the + // pre-checkpoint value. + assertThat(b1.refCnt()).isGreaterThanOrEqualTo(before); + } + + @Test + void testCheckpointStartedRemovesSentinel() throws Exception { + SingleInputGate inputGate = new SingleInputGateBuilder().build(); + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + LocalInputChannel channel = newPushOnlyLocalChannel(inputGate, stateWriter); + inputGate.setInputChannels(channel); + + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(1)); + channel.onRecoveredStateBuffer( + EventSerializer.toBuffer(new RecoveryCheckpointBarrier(1L), false)); + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(2)); + + CheckpointOptions options = + CheckpointOptions.unaligned(CheckpointType.CHECKPOINT, getDefault()); + stateWriter.start(1L, options); + channel.checkpointStarted(new CheckpointBarrier(1L, 0L, options)); + + Optional<InputChannel.BufferAndAvailability> h1 = channel.getNextBuffer(); + assertThat(h1).isPresent(); + Optional<InputChannel.BufferAndAvailability> h2 = channel.getNextBuffer(); + assertThat(h2).isPresent(); + assertThat(h2.get().buffer().getSize()).isEqualTo(2); + } + + @Test + void testCheckpointStartedNestedCpIds() throws Exception { + SingleInputGate inputGate = new SingleInputGateBuilder().build(); + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + LocalInputChannel channel = newPushOnlyLocalChannel(inputGate, stateWriter); + inputGate.setInputChannels(channel); + + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(1)); + channel.onRecoveredStateBuffer( + EventSerializer.toBuffer(new RecoveryCheckpointBarrier(1L), false)); + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(2)); + channel.onRecoveredStateBuffer( + EventSerializer.toBuffer(new RecoveryCheckpointBarrier(2L), false)); + + CheckpointOptions options = + CheckpointOptions.unaligned(CheckpointType.CHECKPOINT, getDefault()); + stateWriter.start(1L, options); + channel.checkpointStarted(new CheckpointBarrier(1L, 0L, options)); + + List<Buffer> persisted = stateWriter.getAddedInput().get(channel.getChannelInfo()); + assertThat(persisted).hasSize(1); + assertThat(persisted.get(0).getSize()).isEqualTo(1); + } + + @Test + void testCheckpointStartedNotInRecoveryUsesMasterPath() throws Exception { + SingleInputGate inputGate = new SingleInputGateBuilder().build(); + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + CountDownLatch upstreamReady = new CountDownLatch(1); + LocalInputChannel channel = newPushOnlyLocalChannel(inputGate, stateWriter, upstreamReady); + inputGate.setInputChannels(channel); + // Channel starts in-recovery; finish appends the sentinel and consuming it flips the + // channel to not-in-recovery so checkpointStarted exercises the master path instead of the + // recovery branch. + upstreamReady.countDown(); + channel.finishRecoveredBufferDelivery(); + channel.getNextBuffer(); + channel.onRecoveredStateConsumed(); + + CheckpointOptions options = + CheckpointOptions.unaligned(CheckpointType.CHECKPOINT, getDefault()); + stateWriter.start(1L, options); + channel.checkpointStarted(new CheckpointBarrier(1L, 0L, options)); + + assertThat(stateWriter.getAddedInput().get(channel.getChannelInfo())).isNullOrEmpty(); + } + + @Test + void testReceivedBuffersHasNoLiveDataBufferIsTrueOnLocal() throws Exception { + // Local has no receivedBuffers; the helper is trivially true. We exercise the + // in-recovery checkpointStarted branch without any "live data" infrastructure to confirm + // the channel does not crash. + SingleInputGate inputGate = new SingleInputGateBuilder().build(); + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + LocalInputChannel channel = newPushOnlyLocalChannel(inputGate, stateWriter); + inputGate.setInputChannels(channel); + + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(1)); + channel.onRecoveredStateBuffer( + EventSerializer.toBuffer(new RecoveryCheckpointBarrier(1L), false)); + + CheckpointOptions options = + CheckpointOptions.unaligned(CheckpointType.CHECKPOINT, getDefault()); + stateWriter.start(1L, options); + channel.checkpointStarted(new CheckpointBarrier(1L, 0L, options)); + } + // --------------------------------------------------------------------------------------------- /** Returns the configured number of buffers for each channel in a random order. */ diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/io/benchmark/SingleInputGateBenchmarkFactory.java b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/io/benchmark/SingleInputGateBenchmarkFactory.java index b850a7cc553..cc0d71c6da4 100644 --- a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/io/benchmark/SingleInputGateBenchmarkFactory.java +++ b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/io/benchmark/SingleInputGateBenchmarkFactory.java @@ -130,7 +130,8 @@ public class SingleInputGateBenchmarkFactory extends SingleInputGateFactory { metrics.getNumBytesInLocalCounter(), metrics.getNumBuffersInLocalCounter(), ChannelStateWriter.NO_OP, - new ArrayDeque<>()); + 0, + false); } @Override
