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 9cb272efdcb82d011347ca438ec347cacc5c349f Author: Rui Fan <[email protected]> AuthorDate: Mon Jul 6 00:26:29 2026 +0200 [FLINK-39521][network] RemoteInputChannel: push-based recovery state Teach RemoteInputChannel to be created directly in a recovery state and to receive recovered buffers through the RecoverableInputChannel interface. The not-in-recovery checkpointStarted branch stays byte-identical to the previous behavior, and with needsRecovery=false the channel behaves exactly as before. - Constructor: drop ArrayDeque<Buffer> initialRecoveredBuffers, add needsRecovery (pre-completed stateConsumedFuture when false, same as LocalInputChannel). The BufferManager starts with credit notification disabled while in recovery (notifyInitiallyEnabled=!needsRecovery). - appendRecoveredBuffer() appends straight into receivedBuffers (NONE subpartition id, recovery sequence numbers from Integer.MIN_VALUE), so the consume path stays identical to the normal case. - recoveryEventStash (@GuardedBy receivedBuffers): ordinary upstream events arriving while credit is suppressed are stashed and replayed after the sentinel is consumed; live data buffers during recovery are a protocol violation asserted in onBuffer. - upstreamReady as CountDownLatch(1), counted down by the first onBuffer or by release; finishRecoveredBufferDelivery() awaits it and appends the EndOfFetchedChannelStateEvent sentinel. - onRecoveredStateConsumed(): unstash + bufferManager.enableNotify() (credit reopens) + complete stateConsumedFuture. - checkReadability() allows in-recovery reads before the partition request client is initialized. - checkpointStarted() in-recovery branch persists the buffers collected by collectPreRecoveryBarrier (walks receivedBuffers past the priority head to the matching RecoveryCheckpointBarrier, retains pre-barrier data buffers, removes the sentinel); only startPersisting is called, so the persist window stays open. A missing barrier declines the checkpoint (CHECKPOINT_DECLINED). - New public needsRecovery() getter, consumed by the netty stack in a later commit of this PR. Mechanical constructor-caller adaptations ride along (UnknownInputChannel, RemoteRecoveredInputChannel, InputChannelBuilder, benchmark factory, netty test channels), all passing needsRecovery=false. RemoteRecoveredInputChannel keeps calling remoteInputChannel.setup() until the conversion rework in the next commit moves setup() into RecoveredInputChannel#toInputChannel. --- .../partition/consumer/RemoteInputChannel.java | 430 ++++++++++++++--- .../consumer/RemoteRecoveredInputChannel.java | 4 +- .../partition/consumer/UnknownInputChannel.java | 33 +- ...editBasedPartitionRequestClientHandlerTest.java | 3 +- .../netty/PartitionRequestRegistrationTest.java | 3 +- .../partition/consumer/InputChannelBuilder.java | 3 +- .../partition/consumer/RemoteInputChannelTest.java | 518 ++++++++++++++++++++- .../benchmark/SingleInputGateBenchmarkFactory.java | 3 +- 8 files changed, 909 insertions(+), 88 deletions(-) diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannel.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannel.java index 001fd45a58a..8b9241e19da 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannel.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannel.java @@ -25,6 +25,7 @@ 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; @@ -46,7 +47,6 @@ import org.apache.flink.runtime.io.network.partition.PrioritizedDeque; import org.apache.flink.runtime.io.network.partition.ResultPartitionID; import org.apache.flink.runtime.io.network.partition.ResultSubpartitionIndexSet; import org.apache.flink.util.ExceptionUtils; -import org.apache.flink.util.concurrent.FutureUtils; import org.apache.flink.shaded.guava33.com.google.common.collect.Iterators; @@ -64,9 +64,11 @@ import java.util.List; import java.util.Optional; import java.util.OptionalLong; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; +import java.util.stream.Stream; import static org.apache.flink.runtime.io.network.buffer.Buffer.DataType.RECOVERY_METADATA; import static org.apache.flink.util.Preconditions.checkArgument; @@ -74,7 +76,7 @@ import static org.apache.flink.util.Preconditions.checkNotNull; import static org.apache.flink.util.Preconditions.checkState; /** An input channel, which requests a remote partition queue. */ -public class RemoteInputChannel extends InputChannel { +public class RemoteInputChannel extends InputChannel implements RecoverableInputChannel { private static final Logger LOG = LoggerFactory.getLogger(RemoteInputChannel.class); private static final int NONE = -1; @@ -109,6 +111,8 @@ public class RemoteInputChannel extends InputChannel { /** The initial number of exclusive buffers assigned to this channel. */ private final int initialCredit; + private final boolean needsRecovery; + /** The milliseconds timeout for partition request listener in result partition manager. */ private final int partitionRequestListenerTimeout; @@ -125,6 +129,45 @@ public class RemoteInputChannel extends InputChannel { private final ChannelStatePersister channelStatePersister; + /** + * Whether the channel is still replaying recovered state. Recovered buffers delivered by the + * spill drain are appended directly to {@link #receivedBuffers}, so the consume path needs no + * recovery-specific branch. 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 that the drain appended after the last recovered + * buffer (see {@link #onRecoveredStateConsumed()}). + */ + @GuardedBy("receivedBuffers") + 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}. + */ + @GuardedBy("receivedBuffers") + private int recoverySequenceNumber = Integer.MIN_VALUE; + + /** + * Ordinary (non-priority) upstream events received while recovery is still in progress. They + * cannot enter {@link #receivedBuffers} ahead of the recovered buffers, so they are stashed + * here and appended once recovery delivery finishes. Credit is suppressed during recovery, so + * the upstream can only send events (never data buffers) before {@link + * #finishRecoveredBufferDelivery()}. + */ + @GuardedBy("receivedBuffers") + private final ArrayDeque<SequenceBuffer> recoveryEventStash = new ArrayDeque<>(); + + /** + * One-shot latch that opens once the upstream reader is registered and the connection is live + * (signalled by the first {@link #onBuffer} or by {@link #releaseAllResources()}). + * Recovery-side awaiters block on it before handing off; once open, {@link + * CountDownLatch#countDown()} on the hot path is a cheap idempotent no-op, unlike completing a + * {@code CompletableFuture}. + */ + private final CountDownLatch upstreamReady; + private long totalQueueSizeInBytes; public RemoteInputChannel( @@ -141,7 +184,42 @@ public class RemoteInputChannel extends InputChannel { Counter numBytesIn, Counter numBuffersIn, ChannelStateWriter stateWriter, - ArrayDeque<Buffer> initialRecoveredBuffers) { + boolean needsRecovery) { + this( + inputGate, + channelIndex, + partitionId, + consumedSubpartitionIndexSet, + connectionId, + connectionManager, + initialBackOff, + maxBackoff, + partitionRequestListenerTimeout, + networkBuffersPerChannel, + numBytesIn, + numBuffersIn, + stateWriter, + needsRecovery, + new CountDownLatch(1)); + } + + @VisibleForTesting + RemoteInputChannel( + SingleInputGate inputGate, + int channelIndex, + ResultPartitionID partitionId, + ResultSubpartitionIndexSet consumedSubpartitionIndexSet, + ConnectionID connectionId, + ConnectionManager connectionManager, + int initialBackOff, + int maxBackoff, + int partitionRequestListenerTimeout, + int networkBuffersPerChannel, + Counter numBytesIn, + Counter numBuffersIn, + ChannelStateWriter stateWriter, + boolean needsRecovery, + CountDownLatch upstreamReady) { super( inputGate, @@ -158,31 +236,15 @@ public class RemoteInputChannel extends InputChannel { this.initialCredit = networkBuffersPerChannel; this.connectionId = checkNotNull(connectionId); this.connectionManager = checkNotNull(connectionManager); + this.needsRecovery = needsRecovery; this.bufferManager = - new BufferManager(inputGate.getMemorySegmentProvider(), this, 0, true); - 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; - for (Buffer buffer : initialRecoveredBuffers) { - // subpartitionId is set to 0 for recovered buffers. This is correct because: - // 1) For single-subpartition channels, the only valid subpartition is 0. - // 2) For multi-subpartition channels (consumedSubpartitionIndexSet.size() > 1), - // RecoveryMetadata events embedded in the recovered buffer sequence track - // the actual subpartition context for proper routing. - SequenceBuffer sequenceBuffer = new SequenceBuffer(buffer, seqNum++, 0); - receivedBuffers.add(sequenceBuffer); - totalQueueSizeInBytes += buffer.getSize(); - } - checkState( - receivedBuffers.size() == expectedCount, - "Buffer migration failed: expected %s buffers but got %s", - expectedCount, - receivedBuffers.size()); + new BufferManager(inputGate.getMemorySegmentProvider(), this, 0, !needsRecovery); + this.channelStatePersister = + new ChannelStatePersister(checkNotNull(stateWriter), getChannelInfo()); + this.inRecovery = needsRecovery; + this.upstreamReady = checkNotNull(upstreamReady); + if (!needsRecovery) { + stateConsumedFuture.complete(null); } } @@ -204,6 +266,117 @@ public class RemoteInputChannel extends InputChannel { bufferManager.requestExclusiveBuffers(initialCredit); } + // ------------------------------------------------------------------------ + // RecoverableInputChannel implementation + // ------------------------------------------------------------------------ + + @Override + public void onRecoveredStateBuffer(Buffer buffer) { + boolean wasEmpty; + synchronized (receivedBuffers) { + if (isReleased.get()) { + buffer.recycleBuffer(); + return; + } + // Migrate recovered buffers from RecoveredInputChannel. These buffers have been + // filtered but not yet consumed by the Task. They are appended to receivedBuffers so + // the consume path stays identical to the non-recovery case. + wasEmpty = appendRecoveredBuffer(buffer); + } + if (wasEmpty) { + notifyChannelNonEmpty(); + } + } + + @Override + public void finishRecoveredBufferDelivery() throws IOException, InterruptedException { + upstreamReady.await(); + boolean wasEmpty; + synchronized (receivedBuffers) { + // A release may have opened the latch instead of the first buffer; bail out so we never + // append to a queue that releaseAllResources() already cleared. + if (isReleased.get()) { + return; + } + checkState(inRecovery, "Recovery delivery already finished."); + // Append the sentinel after the last recovered buffer. The consume path flips out of + // recovery (unstash + reopen credit) only once it polls this sentinel, guaranteeing all + // recovered buffers are consumed first. + wasEmpty = + appendRecoveredBuffer( + EventSerializer.toBuffer( + EndOfFetchedChannelStateEvent.INSTANCE, false)); + } + if (wasEmpty) { + notifyChannelNonEmpty(); + } + } + + /** + * Flips out of recovery once the consume path polls the {@code EndOfFetchedChannelStateEvent} + * sentinel: releases the upstream events stashed during recovery so they are consumed after the + * recovered buffers, then reopens the suppressed credit notifications. + */ + @Override + public void onRecoveredStateConsumed() throws IOException { + synchronized (receivedBuffers) { + checkState(inRecovery, "Recovery already finished."); + inRecovery = false; + recoveryEventStash.forEach(receivedBuffers::add); + recoveryEventStash.clear(); + } + notifyChannelNonEmpty(); + // Credit notifications are suppressed while recovery borrows the exclusive buffers. + bufferManager.enableNotify(); + stateConsumedFuture.complete(null); + } + + @Override + public CompletableFuture<Void> getStateConsumedFuture() { + return stateConsumedFuture; + } + + @Override + public Buffer requestRecoveryBufferBlocking() throws InterruptedException, IOException { + upstreamReady.await(); + // If a release opened the latch instead of the first buffer, requestBufferBlocking() + // detects the released channel and throws CancelTaskException. + return bufferManager.requestBufferBlocking(); + } + + @Override + public void insertRecoveryCheckpointBarrierIfInRecovery(long checkpointId) throws IOException { + boolean wasEmpty = false; + synchronized (receivedBuffers) { + if (!isReleased.get() && inRecovery) { + wasEmpty = + appendRecoveredBuffer( + EventSerializer.toBuffer( + new RecoveryCheckpointBarrier(checkpointId), false)); + } + } + if (wasEmpty) { + notifyChannelNonEmpty(); + } + } + + /** + * Appends a recovered buffer (or {@code RecoveryCheckpointBarrier} sentinel) to {@link + * #receivedBuffers} with a recovery sequence number. + * + * @return {@code true} iff {@code receivedBuffers} transitioned from empty to non-empty. + */ + @GuardedBy("receivedBuffers") + private boolean appendRecoveredBuffer(Buffer buffer) { + boolean wasEmpty = receivedBuffers.isEmpty(); + // Recovered buffers carry no per-buffer subpartition id (NONE): they are snapshotted via + // the recovery path, never via getInflightBuffersUnsafe which is the only consumer of that + // field. + receivedBuffers.add(new SequenceBuffer(buffer, recoverySequenceNumber++, NONE)); + totalQueueSizeInBytes += buffer.getSize(); + return wasEmpty; + } + // ------------------------------------------------------------------------ // Consume // ------------------------------------------------------------------------ @@ -346,14 +519,21 @@ public class RemoteInputChannel extends InputChannel { @Override void releaseAllResources() throws IOException { if (isReleased.compareAndSet(false, 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.")); final ArrayDeque<Buffer> releasedBuffers; synchronized (receivedBuffers) { releasedBuffers = - receivedBuffers.stream() + Stream.concat(receivedBuffers.stream(), recoveryEventStash.stream()) .map(sb -> sb.buffer) .collect(Collectors.toCollection(ArrayDeque::new)); receivedBuffers.clear(); + recoveryEventStash.clear(); } bufferManager.releaseAllBuffers(releasedBuffers); @@ -561,6 +741,10 @@ public class RemoteInputChannel extends InputChannel { return initialCredit; } + public boolean needsRecovery() { + return needsRecovery; + } + public BufferProvider getBufferProvider() throws IOException { if (isReleased.get()) { return null; @@ -599,6 +783,11 @@ public class RemoteInputChannel extends InputChannel { throws IOException { boolean recycleBuffer = true; + // The first buffer from the producer proves the upstream reader is registered and the + // connection is live; release any recovery-side awaiter. On later buffers this is a cheap + // idempotent no-op (the latch count is already zero). + upstreamReady.countDown(); + try { if (expectedSequenceNumber != sequenceNumber) { onError(new BufferReorderingException(expectedSequenceNumber, sequenceNumber)); @@ -636,7 +825,20 @@ public class RemoteInputChannel extends InputChannel { firstPriorityEvent = addPriorityBuffer(sequenceBuffer); recycleBuffer = false; } else { - receivedBuffers.add(sequenceBuffer); + if (inRecovery) { + // The upstream has no credit until recovery delivery finishes, so it can + // only + // send events here, never data buffers. Stash ordinary events so they are + // consumed after the recovered buffers; data buffers are a protocol + // violation. + checkState( + !buffer.isBuffer(), + "Received live data buffer during recovery on channel %s", + getChannelInfo()); + recoveryEventStash.add(sequenceBuffer); + } else { + receivedBuffers.add(sequenceBuffer); + } recycleBuffer = false; if (dataType.requiresAnnouncement()) { firstPriorityEvent = addPriorityBuffer(announce(sequenceBuffer)); @@ -716,35 +918,153 @@ public class RemoteInputChannel extends InputChannel { } /** - * Spills all queued buffers on checkpoint start. If barrier has already been received (and - * reordered), spill only the overtaken buffers. + * Persists inflight data on checkpoint start. During recovery, persists recovered buffers + * before the matching RecoveryCheckpointBarrier sentinel; after recovery, uses the normal + * remote-channel barrier sequence tracking and persists overtaken live buffers. */ public void checkpointStarted(CheckpointBarrier barrier) throws CheckpointException { - synchronized (receivedBuffers) { - if (barrier.getId() < lastBarrierId) { - throw new CheckpointException( - String.format( - "Sequence number for checkpoint %d is not known (it was likely been overwritten by a newer checkpoint %d)", - barrier.getId(), lastBarrierId), - CheckpointFailureReason - .CHECKPOINT_SUBSUMED); // currently, at most one active unaligned - // checkpoint is possible - } else if (barrier.getId() > lastBarrierId) { - // This channel has received some obsolete barrier, older compared to the - // checkpointId - // which we are processing right now, and we should ignore that obsoleted checkpoint - // barrier sequence number. - resetLastBarrier(); + try { + List<Buffer> toPersist; + synchronized (receivedBuffers) { + if (inRecovery) { + toPersist = collectPreRecoveryBarrier(barrier.getId()); + } else { + if (barrier.getId() < lastBarrierId) { + // Currently, at most one active unaligned checkpoint is possible. + throw new CheckpointException( + String.format( + "Sequence number for checkpoint %d is not known (it was likely been overwritten by a newer checkpoint %d)", + barrier.getId(), lastBarrierId), + CheckpointFailureReason.CHECKPOINT_SUBSUMED); + } else if (barrier.getId() > lastBarrierId) { + // This channel has received some obsolete barrier, older compared to the + // checkpointId which we are processing right now, and we should ignore that + // obsoleted checkpoint barrier sequence number. + resetLastBarrier(); + } + toPersist = getInflightBuffersUnsafe(barrier.getId()); + } + channelStatePersister.startPersisting(barrier.getId(), toPersist); } + } catch (IOException e) { + throw new CheckpointException( + "Failed to extract recovered buffers for checkpoint " + barrier.getId(), + CheckpointFailureReason.CHECKPOINT_DECLINED, + e); + } + } - channelStatePersister.startPersisting( - barrier.getId(), getInflightBuffersUnsafe(barrier.getId())); + /** + * Walks {@link #receivedBuffers} (skipping priority events) up to the {@link + * RecoveryCheckpointBarrier} sentinel matching {@code checkpointId}, retaining each pre-barrier + * recovered data buffer and removing the sentinel. During recovery the upstream has no credit, + * so {@code receivedBuffers} holds only recovered buffers, sentinels, and priority events — no + * live data buffers. + * + * <p>The scan runs on any checkpoint that starts while the channel is still {@code inRecovery}, + * whether or not the drain has finished delivering. It must skip the {@code + * EndOfFetchedChannelStateEvent} because that sentinel can sit ahead of the {@code + * RecoveryCheckpointBarrier}: once the drain finishes it appends {@code + * EndOfFetchedChannelStateEvent}, and if a checkpoint is then triggered before the consume path + * polls it, {@link #insertRecoveryCheckpointBarrierIfInRecovery} appends the {@code + * RecoveryCheckpointBarrier} behind it. The incoming {@link CheckpointBarrier} is a priority + * event kept at the head, skipped by advancing past the priority region. + * + * @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). + */ + @GuardedBy("receivedBuffers") + private List<Buffer> collectPreRecoveryBarrier(long checkpointId) throws IOException { + assert Thread.holdsLock(receivedBuffers); + List<Buffer> retained = new ArrayList<>(); + List<SequenceBuffer> subsumed = new ArrayList<>(); + SequenceBuffer sentinel = null; + try { + Iterator<SequenceBuffer> it = receivedBuffers.iterator(); + // Priority events are stored separately at the head and never carry recovered data. + Iterators.advance(it, receivedBuffers.getNumPriorityElements()); + while (it.hasNext()) { + SequenceBuffer sb = it.next(); + RecoveryCheckpointBarrier barrier = asRecoveryCheckpointBarrier(sb.buffer); + if (barrier != null) { + long barrierId = barrier.getCheckpointId(); + if (barrierId == checkpointId) { + sentinel = sb; + break; + } + if (barrierId > checkpointId) { + throw new IOException( + "Found RecoveryCheckpointBarrier for a later checkpoint " + + barrierId + + " before the target checkpoint " + + checkpointId + + " in receivedBuffers 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()); + subsumed.add(sb); + continue; + } + // Skip non-data events (e.g. the EndOfFetchedChannelStateEvent sentinel appended + // after the recovered buffers): only recovered data buffers are snapshotted. + if (sb.buffer.isBuffer()) { + retained.add(sb.buffer.retainBuffer()); + } + } + } catch (IOException e) { + releaseRetainedBuffers(retained); + throw e; + } + if (sentinel == null) { + releaseRetainedBuffers(retained); + throw new IOException( + "Missing RecoveryCheckpointBarrier for checkpoint " + + checkpointId + + " in receivedBuffers for channel " + + getChannelInfo()); } + // receivedBuffers is a PrioritizedDeque whose iterator() is read-only; remove matched and + // subsumed sentinels by identity through its mutable removal API. + for (SequenceBuffer s : subsumed) { + removeRecoverySentinel(s); + } + removeRecoverySentinel(sentinel); + return retained; } - @Override - public CompletableFuture<Void> getStateConsumedFuture() { - return FutureUtils.completedVoidFuture(); + @GuardedBy("receivedBuffers") + private void removeRecoverySentinel(SequenceBuffer sentinel) { + receivedBuffers.getAndRemove(sb -> sb == sentinel); + totalQueueSizeInBytes -= sentinel.buffer.getSize(); + sentinel.buffer.recycleBuffer(); + } + + 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; } public void checkpointStopped(long checkpointId) { @@ -917,13 +1237,13 @@ public class RemoteInputChannel extends InputChannel { } /** - * When receivedBuffers contains migrated buffers from RecoveredInputChannel, they can be read - * before requestSubpartitions(). In that case only check for errors. Once migrated buffers are - * drained, require full client initialization check. + * Allows reads while recovery data or already queued network data is available before the + * remote partition request is fully initialized. If neither recovery nor queued data can + * satisfy the read, require the partition request client to be initialized. */ private void checkReadability() throws IOException { assert Thread.holdsLock(receivedBuffers); - if (receivedBuffers.isEmpty()) { + if (!inRecovery && receivedBuffers.isEmpty()) { checkPartitionRequestQueueInitialized(); } else { checkError(); diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteRecoveredInputChannel.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteRecoveredInputChannel.java index 2cfff6f5e79..6b0279ad3e2 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteRecoveredInputChannel.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteRecoveredInputChannel.java @@ -70,6 +70,8 @@ public class RemoteRecoveredInputChannel extends RecoveredInputChannel { @Override protected InputChannel toInputChannelInternal(ArrayDeque<Buffer> remainingBuffers) throws IOException { + // 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. RemoteInputChannel remoteInputChannel = new RemoteInputChannel( inputGate, @@ -85,7 +87,7 @@ public class RemoteRecoveredInputChannel extends RecoveredInputChannel { numBytesIn, numBuffersIn, channelStateWriter, - remainingBuffers); + false); remoteInputChannel.setup(); return remoteInputChannel; } 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 78d6f024fd4..9db9df9c613 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 @@ -36,7 +36,6 @@ import org.apache.flink.util.concurrent.FutureUtils; import javax.annotation.Nullable; import java.io.IOException; -import java.util.ArrayDeque; import java.util.Optional; import java.util.concurrent.CompletableFuture; @@ -178,21 +177,23 @@ class UnknownInputChannel extends InputChannel implements ChannelStateHolder { public RemoteInputChannel toRemoteInputChannel( ConnectionID producerAddress, ResultPartitionID resultPartitionID) { - return new RemoteInputChannel( - inputGate, - getChannelIndex(), - resultPartitionID, - consumedSubpartitionIndexSet, - checkNotNull(producerAddress), - connectionManager, - initialBackoff, - maxBackoff, - partitionRequestListenerTimeout, - networkBuffersPerChannel, - metrics.getNumBytesInRemoteCounter(), - metrics.getNumBuffersInRemoteCounter(), - channelStateWriter == null ? ChannelStateWriter.NO_OP : channelStateWriter, - new ArrayDeque<>()); + RemoteInputChannel channel = + new RemoteInputChannel( + inputGate, + getChannelIndex(), + resultPartitionID, + consumedSubpartitionIndexSet, + checkNotNull(producerAddress), + connectionManager, + initialBackoff, + maxBackoff, + partitionRequestListenerTimeout, + networkBuffersPerChannel, + metrics.getNumBytesInRemoteCounter(), + metrics.getNumBuffersInRemoteCounter(), + channelStateWriter == null ? ChannelStateWriter.NO_OP : channelStateWriter, + false); + return channel; } public LocalInputChannel toLocalInputChannel(ResultPartitionID resultPartitionID) { diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/netty/CreditBasedPartitionRequestClientHandlerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/netty/CreditBasedPartitionRequestClientHandlerTest.java index d4b162304b8..d5801803696 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/netty/CreditBasedPartitionRequestClientHandlerTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/netty/CreditBasedPartitionRequestClientHandlerTest.java @@ -68,7 +68,6 @@ import org.junit.jupiter.params.provider.MethodSource; import java.io.IOException; import java.net.InetSocketAddress; -import java.util.ArrayDeque; import java.util.stream.Stream; import static org.apache.flink.runtime.io.network.netty.PartitionRequestQueueTest.blockChannel; @@ -956,7 +955,7 @@ class CreditBasedPartitionRequestClientHandlerTest { new SimpleCounter(), new SimpleCounter(), ChannelStateWriter.NO_OP, - new ArrayDeque<>()); + false); this.expectedMessage = expectedMessage; } diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/netty/PartitionRequestRegistrationTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/netty/PartitionRequestRegistrationTest.java index e3cfb55e340..4080858b74b 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/netty/PartitionRequestRegistrationTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/netty/PartitionRequestRegistrationTest.java @@ -44,7 +44,6 @@ import org.apache.flink.shaded.netty4.io.netty.channel.Channel; import org.junit.jupiter.api.Test; -import java.util.ArrayDeque; import java.util.Optional; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; @@ -250,7 +249,7 @@ class PartitionRequestRegistrationTest { new SimpleCounter(), new SimpleCounter(), ChannelStateWriter.NO_OP, - new ArrayDeque<>()); + false); this.latch = latch; } 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 1130c13bd4f..e69f8414b43 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 @@ -198,7 +198,8 @@ public class InputChannelBuilder { metrics.getNumBytesInRemoteCounter(), metrics.getNumBuffersInRemoteCounter(), stateWriter, - new ArrayDeque<>()); + needsRecovery, + upstreamReady); } public LocalRecoveredInputChannel buildLocalRecoveredChannel(SingleInputGate inputGate) { diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannelTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannelTest.java index e47de93c9e8..9b1ccaabdb8 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannelTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannelTest.java @@ -23,10 +23,14 @@ import org.apache.flink.core.memory.MemorySegmentFactory; import org.apache.flink.core.testutils.OneShotLatch; 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.InputChannelInfo; +import org.apache.flink.runtime.checkpoint.channel.RecordingChannelStateWriter; +import org.apache.flink.runtime.checkpoint.channel.RecoveryCheckpointBarrier; +import org.apache.flink.runtime.clusterframework.types.ResourceID; import org.apache.flink.runtime.execution.CancelTaskException; import org.apache.flink.runtime.execution.ExecutionState; import org.apache.flink.runtime.io.network.ConnectionID; @@ -37,6 +41,7 @@ import org.apache.flink.runtime.io.network.PartitionRequestClient; import org.apache.flink.runtime.io.network.TestingConnectionManager; import org.apache.flink.runtime.io.network.TestingPartitionRequestClient; import org.apache.flink.runtime.io.network.api.CheckpointBarrier; +import org.apache.flink.runtime.io.network.api.EndOfPartitionEvent; 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.Buffer.DataType; @@ -72,6 +77,7 @@ import org.junit.jupiter.api.extension.RegisterExtension; import javax.annotation.Nullable; import java.io.IOException; +import java.net.InetSocketAddress; import java.util.ArrayDeque; import java.util.ArrayList; import java.util.List; @@ -80,6 +86,7 @@ import java.util.Queue; import java.util.Random; 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.Future; @@ -2079,15 +2086,9 @@ class RemoteInputChannelTest { // given: RemoteInputChannel 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)); - ConnectionID connectionId = - new ConnectionID( - org.apache.flink.runtime.clusterframework.types.ResourceID.generate(), - new java.net.InetSocketAddress("localhost", 0), - 0); + new ConnectionID(ResourceID.generate(), new InetSocketAddress("localhost", 0), 0); + CountDownLatch upstreamReady = new CountDownLatch(1); RemoteInputChannel channel = new RemoteInputChannel( inputGate, @@ -2104,10 +2105,16 @@ class RemoteInputChannelTest { new SimpleCounter(), new SimpleCounter(), ChannelStateWriter.NO_OP, - recoveredBuffers); + 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<BufferAndAvailability> first = channel.getNextBuffer(); assertThat(first).isPresent(); @@ -2119,6 +2126,499 @@ class RemoteInputChannelTest { assertThat(second.get().buffer().getSize()).isEqualTo(20); } + // --------------------------------------------------------------------------------------------- + // RecoverableInputChannel push-based recovery tests + // --------------------------------------------------------------------------------------------- + + @Test + void testOnRecoveredStateBufferEnqueues() throws Exception { + SingleInputGate inputGate = createSingleInputGate(1); + RemoteInputChannel channel = + InputChannelBuilder.newBuilder() + .setStateWriter(ChannelStateWriter.NO_OP) + .setNeedsRecovery(true) + .buildRemoteChannel(inputGate); + inputGate.setInputChannels(channel); + + Buffer b1 = TestBufferFactory.createBuffer(11); + Buffer b2 = TestBufferFactory.createBuffer(22); + channel.onRecoveredStateBuffer(b1); + channel.onRecoveredStateBuffer(b2); + + Optional<BufferAndAvailability> first = channel.getNextBuffer(); + assertThat(first).isPresent(); + assertThat(first.get().buffer().getSize()).isEqualTo(11); + Optional<BufferAndAvailability> second = channel.getNextBuffer(); + assertThat(second).isPresent(); + assertThat(second.get().buffer().getSize()).isEqualTo(22); + } + + @Test + void testRecoveredBuffersConsumedBeforeStashedEventsThenSentinelFlipsRecovery() + throws Exception { + SingleInputGate inputGate = createSingleInputGate(1); + CountDownLatch upstreamReady = new CountDownLatch(1); + RemoteInputChannel channel = + InputChannelBuilder.newBuilder() + .setStateWriter(ChannelStateWriter.NO_OP) + .setNeedsRecovery(true) + .setUpstreamReady(upstreamReady) + .buildRemoteChannel(inputGate); + inputGate.setInputChannels(channel); + + // Recovered buffer arrives via the drain; an ordinary upstream event arrives via onBuffer + // while still in recovery and must be stashed (events carry no credit). + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(11)); + // backlog=-1: events carry no backlog, and this channel has no floating-buffer pool wired. + channel.onBuffer(EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE, false), 0, -1, 0); + upstreamReady.countDown(); + channel.finishRecoveredBufferDelivery(); + + // Recovered buffer is consumed first. + Optional<BufferAndAvailability> recovered = channel.getNextBuffer(); + assertThat(recovered).isPresent(); + assertThat(recovered.get().buffer().getSize()).isEqualTo(11); + + // Then the sentinel; the stashed event is not yet visible (still in recovery). + Optional<BufferAndAvailability> sentinel = channel.getNextBuffer(); + assertThat(sentinel).isPresent(); + assertThat(EventSerializer.fromBuffer(sentinel.get().buffer(), getClass().getClassLoader())) + .isInstanceOf(EndOfFetchedChannelStateEvent.class); + + // The gate consumes the sentinel externally: flips out of recovery and unstashes the event. + channel.onRecoveredStateConsumed(); + Optional<BufferAndAvailability> stashed = channel.getNextBuffer(); + assertThat(stashed).isPresent(); + assertThat(EventSerializer.fromBuffer(stashed.get().buffer(), getClass().getClassLoader())) + .isInstanceOf(EndOfPartitionEvent.class); + } + + @Test + void testOnBufferRejectsLiveDataBufferDuringRecovery() throws Exception { + SingleInputGate inputGate = createSingleInputGate(1); + RemoteInputChannel channel = + InputChannelBuilder.newBuilder() + .setStateWriter(ChannelStateWriter.NO_OP) + .setNeedsRecovery(true) + .buildRemoteChannel(inputGate); + inputGate.setInputChannels(channel); + + assertThatThrownBy(() -> channel.onBuffer(TestBufferFactory.createBuffer(1), 0, 0, 0)) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("Received live data buffer during recovery"); + } + + @Test + void testOnRecoveredStateBufferOnReleasedChannelIsSilentlyRecycled() throws Exception { + SingleInputGate inputGate = createSingleInputGate(1); + RemoteInputChannel channel = + InputChannelBuilder.newBuilder() + .setStateWriter(ChannelStateWriter.NO_OP) + .setNeedsRecovery(true) + .buildRemoteChannel(inputGate); + inputGate.setInputChannels(channel); + channel.releaseAllResources(); + + Buffer b = TestBufferFactory.createBuffer(33); + channel.onRecoveredStateBuffer(b); + + assertThat(b.isRecycled()).isTrue(); + } + + @Test + void testOnRecoveredStateBufferNotifiesChannelNonEmptyOnEmptyToNonEmptyTransition() + throws Exception { + SingleInputGate inputGate = createSingleInputGate(1); + RemoteInputChannel channel = + InputChannelBuilder.newBuilder() + .setStateWriter(ChannelStateWriter.NO_OP) + .setNeedsRecovery(true) + .buildRemoteChannel(inputGate); + inputGate.setInputChannels(channel); + + CompletableFuture<?> availability = inputGate.getAvailableFuture(); + assertThat(availability).isNotDone(); + + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(1)); + assertThat(availability).isDone(); + } + + @Test + void testInRecoveryBoundaryFlagFalseQueueEmptyReturnsEmpty() throws Exception { + SingleInputGate inputGate = createSingleInputGate(1); + RemoteInputChannel channel = + InputChannelBuilder.newBuilder() + .setStateWriter(ChannelStateWriter.NO_OP) + .setNeedsRecovery(true) + .buildRemoteChannel(inputGate); + inputGate.setInputChannels(channel); + + // Force in-recovery + empty queue: push then poll a buffer, then push another while + // delaying finish. After the consumer drains the staged buffer, queue=empty and + // flag=false (no finishRecoveredBufferDelivery called yet). getNextBuffer must return + // empty. + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(1)); + channel.getNextBuffer(); + // Simulate an explicit recovery context where the producer signals "not done yet". + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(2)); + channel.getNextBuffer(); + // The boundary case "flag=false (drain still running) + queue empty" should return empty. + // To set this state explicitly, we deliberately do not call finishReadRecoveredState. + Optional<BufferAndAvailability> result = channel.getNextBuffer(); + assertThat(result).isNotPresent(); + } + + @Test + void testInRecoveryBoundaryFlagFalseQueueNonEmptyPolls() throws Exception { + SingleInputGate inputGate = createSingleInputGate(1); + RemoteInputChannel channel = + InputChannelBuilder.newBuilder() + .setStateWriter(ChannelStateWriter.NO_OP) + .setNeedsRecovery(true) + .buildRemoteChannel(inputGate); + inputGate.setInputChannels(channel); + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(7)); + + Optional<BufferAndAvailability> r = channel.getNextBuffer(); + assertThat(r).isPresent(); + assertThat(r.get().buffer().getSize()).isEqualTo(7); + } + + @Test + void testInRecoveryBoundaryFlagTrueQueueNonEmptyPolls() throws Exception { + SingleInputGate inputGate = createSingleInputGate(1); + CountDownLatch upstreamReady = new CountDownLatch(1); + RemoteInputChannel channel = + InputChannelBuilder.newBuilder() + .setStateWriter(ChannelStateWriter.NO_OP) + .setNeedsRecovery(true) + .setUpstreamReady(upstreamReady) + .buildRemoteChannel(inputGate); + inputGate.setInputChannels(channel); + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(8)); + upstreamReady.countDown(); + channel.finishRecoveredBufferDelivery(); + + Optional<BufferAndAvailability> r = channel.getNextBuffer(); + assertThat(r).isPresent(); + assertThat(r.get().buffer().getSize()).isEqualTo(8); + } + + @Test + void testFinishWithNoRecoveredBuffersEmitsSentinelThenFallsToMasterPath() throws Exception { + // Wire a real network pool so requestSubpartitions() can succeed and the master path can + // poll receivedBuffers. + final NetworkBufferPool networkBufferPool = new NetworkBufferPool(4, 4096); + final CountDownLatch upstreamReady = new CountDownLatch(1); + try { + SingleInputGate inputGate = + new SingleInputGateBuilder() + .setBufferPoolFactory(networkBufferPool.createBufferPool(1, 4)) + .setSegmentProvider(networkBufferPool) + .setChannelFactory( + (builder, gate) -> + builder.setNeedsRecovery(true) + .setUpstreamReady(upstreamReady) + .buildRemoteChannel(gate)) + .build(); + inputGate.setup(); + RemoteInputChannel channel = (RemoteInputChannel) inputGate.getChannel(0); + upstreamReady.countDown(); + // Even with no recovered buffers, finish appends the EndOfFetchedChannelStateEvent + // sentinel so the consume path can flip out of recovery in order. + channel.finishRecoveredBufferDelivery(); + inputGate.requestPartitions(); + + Optional<BufferAndAvailability> sentinel = channel.getNextBuffer(); + assertThat(sentinel).isPresent(); + assertThat( + EventSerializer.fromBuffer( + sentinel.get().buffer(), getClass().getClassLoader())) + .isInstanceOf(EndOfFetchedChannelStateEvent.class); + + // Consuming the sentinel (done externally by the gate) flips the channel out of + // recovery; afterwards the master path is taken and there is no more queued data. + channel.onRecoveredStateConsumed(); + assertThat(channel.getNextBuffer()).isNotPresent(); + } finally { + networkBufferPool.destroy(); + } + } + + @Test + void testMoreAvailableNoneWhenLastRecoveredBufferAndDrainNotFinished() throws Exception { + // While the channel is in recovery and the drain has not finished, the last currently + // queued recovered buffer must report NONE as its next data type: no live data can enter + // receivedBuffers (the upstream has no credit), so there is nothing else to expose yet. + final NetworkBufferPool networkBufferPool = new NetworkBufferPool(4, 4096); + try { + SingleInputGate inputGate = + new SingleInputGateBuilder() + .setBufferPoolFactory(networkBufferPool.createBufferPool(1, 4)) + .setSegmentProvider(networkBufferPool) + .setChannelFactory( + (builder, gate) -> + builder.setNeedsRecovery(true).buildRemoteChannel(gate)) + .build(); + inputGate.setup(); + RemoteInputChannel channel = (RemoteInputChannel) inputGate.getChannel(0); + + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(11)); + + Optional<BufferAndAvailability> recoveredBuf = channel.getNextBuffer(); + assertThat(recoveredBuf).isPresent(); + assertThat(recoveredBuf.get().buffer().getSize()).isEqualTo(11); + // Drain not finished and queue now empty: nothing more is available yet. + assertThat(recoveredBuf.get().moreAvailable()).isFalse(); + } finally { + networkBufferPool.destroy(); + } + } + + @Test + void testPriorityEventDuringRecoveryViaAddPriorityBuffer() throws Exception { + final NetworkBufferPool networkBufferPool = new NetworkBufferPool(4, 4096); + try { + SingleInputGate inputGate = + new SingleInputGateBuilder() + .setBufferPoolFactory(networkBufferPool.createBufferPool(1, 4)) + .setSegmentProvider(networkBufferPool) + .setChannelFactory( + (builder, gate) -> + builder.setNeedsRecovery(true).buildRemoteChannel(gate)) + .build(); + inputGate.setup(); + RemoteInputChannel channel = (RemoteInputChannel) inputGate.getChannel(0); + + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(11)); + + CheckpointBarrier barrier = new CheckpointBarrier(1L, 0L, UNALIGNED); + channel.onBuffer(toBuffer(barrier, true), 0, 0, 0); + + Optional<BufferAndAvailability> first = channel.getNextBuffer(); + assertThat(first).isPresent(); + assertThat(first.get().buffer().getDataType().hasPriority()).isTrue(); + + Optional<BufferAndAvailability> second = channel.getNextBuffer(); + assertThat(second).isPresent(); + assertThat(second.get().buffer().getSize()).isEqualTo(11); + } finally { + networkBufferPool.destroy(); + } + } + + @Test + void testCheckpointStartedScansRecoveredBuffersUpToBarrier() throws Exception { + SingleInputGate inputGate = createSingleInputGate(1); + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + RemoteInputChannel channel = + InputChannelBuilder.newBuilder() + .setStateWriter(stateWriter) + .setNeedsRecovery(true) + .buildRemoteChannel(inputGate); + 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); + + stateWriter.start(1L, UNALIGNED); + channel.checkpointStarted(new CheckpointBarrier(1L, 0L, UNALIGNED)); + + 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 = createSingleInputGate(1); + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + RemoteInputChannel channel = + InputChannelBuilder.newBuilder() + .setStateWriter(stateWriter) + .setNeedsRecovery(true) + .buildRemoteChannel(inputGate); + inputGate.setInputChannels(channel); + + Buffer b1 = TestBufferFactory.createBuffer(1); + channel.onRecoveredStateBuffer(b1); + int refCntBefore = b1.refCnt(); + + stateWriter.start(1L, UNALIGNED); + + // 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, UNALIGNED))) + .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 = createSingleInputGate(1); + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + RemoteInputChannel channel = + InputChannelBuilder.newBuilder() + .setStateWriter(stateWriter) + .setNeedsRecovery(true) + .buildRemoteChannel(inputGate); + inputGate.setInputChannels(channel); + + Buffer b1 = TestBufferFactory.createBuffer(1); + channel.onRecoveredStateBuffer(b1); + channel.onRecoveredStateBuffer( + EventSerializer.toBuffer(new RecoveryCheckpointBarrier(1L), false)); + + int before = b1.refCnt(); + stateWriter.start(1L, UNALIGNED); + channel.checkpointStarted(new CheckpointBarrier(1L, 0L, UNALIGNED)); + // After retainBuffer the buffer remains live for both the queue read and the writer copy. + assertThat(b1.refCnt()).isGreaterThanOrEqualTo(before); + } + + @Test + void testCheckpointStartedRemovesSentinel() throws Exception { + SingleInputGate inputGate = createSingleInputGate(1); + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + RemoteInputChannel channel = + InputChannelBuilder.newBuilder() + .setStateWriter(stateWriter) + .setNeedsRecovery(true) + .buildRemoteChannel(inputGate); + inputGate.setInputChannels(channel); + + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(1)); + channel.onRecoveredStateBuffer( + EventSerializer.toBuffer(new RecoveryCheckpointBarrier(1L), false)); + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(2)); + + stateWriter.start(1L, UNALIGNED); + channel.checkpointStarted(new CheckpointBarrier(1L, 0L, UNALIGNED)); + + Optional<BufferAndAvailability> head = channel.getNextBuffer(); + assertThat(head).isPresent(); + Optional<BufferAndAvailability> nextHead = channel.getNextBuffer(); + assertThat(nextHead).isPresent(); + assertThat(nextHead.get().buffer().getSize()).isEqualTo(2); + } + + @Test + void testCheckpointStartedNestedCpIds() throws Exception { + SingleInputGate inputGate = createSingleInputGate(1); + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + RemoteInputChannel channel = + InputChannelBuilder.newBuilder() + .setStateWriter(stateWriter) + .setNeedsRecovery(true) + .buildRemoteChannel(inputGate); + 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)); + + stateWriter.start(1L, UNALIGNED); + channel.checkpointStarted(new CheckpointBarrier(1L, 0L, UNALIGNED)); + + List<Buffer> persisted1 = stateWriter.getAddedInput().get(channel.getChannelInfo()); + assertThat(persisted1).hasSize(1); + assertThat(persisted1.get(0).getSize()).isEqualTo(1); + } + + @Test + void testCheckpointStartedNotInRecoveryUsesMasterPath() throws Exception { + SingleInputGate inputGate = createSingleInputGate(1); + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + RemoteInputChannel channel = + InputChannelBuilder.newBuilder() + .setStateWriter(stateWriter) + .buildRemoteChannel(inputGate); + inputGate.setInputChannels(channel); + channel.requestSubpartitions(); + stateWriter.start(7L, UNALIGNED); + channel.checkpointStarted(new CheckpointBarrier(7L, 0L, UNALIGNED)); + + assertThat(stateWriter.getAddedInput().get(channel.getChannelInfo())).isNullOrEmpty(); + } + + @Test + void testReceivedBuffersHasNoLiveDataBufferDetectsLiveData() throws Exception { + final NetworkBufferPool networkBufferPool = new NetworkBufferPool(4, 4096); + try { + SingleInputGate inputGate = + new SingleInputGateBuilder() + .setBufferPoolFactory(networkBufferPool.createBufferPool(1, 4)) + .setSegmentProvider(networkBufferPool) + .setChannelFactory( + (builder, gate) -> + builder.setNeedsRecovery(true).buildRemoteChannel(gate)) + .build(); + inputGate.setup(); + RemoteInputChannel channel = (RemoteInputChannel) inputGate.getChannel(0); + + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(1)); + // During recovery the upstream has no credit and can only send events. A live data + // buffer is a protocol violation that onBuffer must reject at the entry point. + assertThatThrownBy(() -> channel.onBuffer(TestBufferFactory.createBuffer(1), 0, 0, 0)) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("Received live data buffer during recovery"); + } finally { + networkBufferPool.destroy(); + } + } + + @Test + void testReceivedBuffersHasNoLiveDataBufferAcceptsPriorityOnly() throws Exception { + final NetworkBufferPool networkBufferPool = new NetworkBufferPool(4, 4096); + try { + RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter(); + SingleInputGate inputGate = + new SingleInputGateBuilder() + .setBufferPoolFactory(networkBufferPool.createBufferPool(1, 4)) + .setSegmentProvider(networkBufferPool) + .setChannelFactory( + (builder, gate) -> + builder.setStateWriter(stateWriter) + .setNeedsRecovery(true) + .buildRemoteChannel(gate)) + .build(); + inputGate.setup(); + RemoteInputChannel channel = (RemoteInputChannel) inputGate.getChannel(0); + + channel.onRecoveredStateBuffer(TestBufferFactory.createBuffer(1)); + // Mirror what snapshotAndInsertBarriers does: push the + // RecoveryCheckpointBarrier sentinel so collectPreRecoveryBarrier finds it. + channel.onRecoveredStateBuffer( + EventSerializer.toBuffer(new RecoveryCheckpointBarrier(1L), false)); + // Priority event in receivedBuffers is OK (!isBuffer()). + CheckpointBarrier priorityBarrier = new CheckpointBarrier(1L, 0L, UNALIGNED); + channel.onBuffer(toBuffer(priorityBarrier, true), 0, 0, 0); + + stateWriter.start(1L, UNALIGNED); + channel.checkpointStarted(new CheckpointBarrier(1L, 0L, UNALIGNED)); + } finally { + networkBufferPool.destroy(); + } + } + private static final class TestBufferPool extends NoOpBufferPool { @Override 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 cc0d71c6da4..b3f0aec9dc5 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 @@ -37,7 +37,6 @@ import org.apache.flink.runtime.shuffle.NettyShuffleDescriptor; import org.apache.flink.runtime.taskmanager.NettyShuffleEnvironmentConfiguration; import java.io.IOException; -import java.util.ArrayDeque; /** * A benchmark-specific input gate factory which overrides the respective methods of creating {@link @@ -187,7 +186,7 @@ public class SingleInputGateBenchmarkFactory extends SingleInputGateFactory { metrics.getNumBytesInRemoteCounter(), metrics.getNumBuffersInRemoteCounter(), ChannelStateWriter.NO_OP, - new ArrayDeque<>()); + false); } @Override
