rkhachatryan commented on code in PR #28652: URL: https://github.com/apache/flink/pull/28652#discussion_r3546635744
########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/EndOfFetchedChannelStateEvent.java: ########## @@ -0,0 +1,75 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.consumer; + +import org.apache.flink.core.memory.DataInputView; +import org.apache.flink.core.memory.DataOutputView; +import org.apache.flink.runtime.event.RuntimeEvent; + +/** + * Marks the tail of recovered buffers that the spill drain pushed into a {@link + * RecoverableInputChannel}. The consume path polls this sentinel to learn the exact moment all + * recovered buffers have been consumed; it is never delivered to the operator. It is distinct from + * {@link EndOfInputChannelStateEvent} (which terminates the {@link RecoveredInputChannel} read + * stream) so the two recovery handoffs cannot be confused. + */ +public class EndOfFetchedChannelStateEvent extends RuntimeEvent { + + /** The singleton instance of this event. */ + public static final EndOfFetchedChannelStateEvent INSTANCE = + new EndOfFetchedChannelStateEvent(); + + // ------------------------------------------------------------------------ + + // not instantiable + private EndOfFetchedChannelStateEvent() {} + + // ------------------------------------------------------------------------ + + @Override + public void write(DataOutputView out) { + throw new UnsupportedOperationException( + "EndOfFetchedChannelStateEvent must be serialized via EventSerializer's dedicated" + + " type-tag path, not reflective write()."); + } + + @Override + public void read(DataInputView in) { + throw new UnsupportedOperationException( + "EndOfFetchedChannelStateEvent must be deserialized via EventSerializer's dedicated" + + " type-tag path, not reflective read()."); + } + + // ------------------------------------------------------------------------ + + @Override + public int hashCode() { + return 20250814; Review Comment: What happened on Aug 14?:) ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/LocalInputChannel.java: ########## @@ -113,48 +159,211 @@ public LocalInputChannel( 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++); - toBeConsumedBuffers.add(bufferAndBacklog); + 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; + if (!needsRecovery) { + stateConsumedFuture.complete(null); + } + } + + @Override + void setup() throws IOException { + if (needsRecovery && networkBuffersPerChannel > 0) { + bufferManager.requestExclusiveBuffers(networkBuffersPerChannel); + } + } + + // ------------------------------------------------------------------------ + // RecoverableInputChannel implementation + // ------------------------------------------------------------------------ + + @Override + public void onRecoveredStateBuffer(Buffer buffer) { + boolean wasEmpty; + synchronized (recoveredBuffers) { + if (isReleased) { + buffer.recycleBuffer(); + return; } - checkState( - toBeConsumedBuffers.size() == expectedCount, - "Buffer migration failed: expected %s buffers but got %s", - expectedCount, - toBeConsumedBuffers.size()); + // 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 { + upstreamReady.join(); + boolean wasEmpty; + synchronized (recoveredBuffers) { + 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.join(); + 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; } + notifyChannelNonEmpty(); + stateConsumedFuture.complete(null); + } + + @Override + public CompletableFuture<Void> getStateConsumedFuture() { + 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. + * + * @throws IOException 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(); + if (isRecoveryCheckpointBarrier(b, checkpointId)) { + it.remove(); + b.recycleBuffer(); + return retained; + } + if (b.isBuffer()) { + retained.add(b.retainBuffer()); Review Comment: What happens if `b` is a `RecoveryCheckpointBarrier` for a *different* checkpoint id than the one we're collecting for? It falls through both branches here (not a data buffer, not the matching sentinel) and is left in the deque without being removed by the iterator. Is that intentional? ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/LocalInputChannel.java: ########## @@ -504,18 +772,33 @@ void releaseAllResources() throws IOException { if (!isReleased) { isReleased = true; + upstreamReady.completeExceptionally(new CancelTaskException("Channel released.")); + + // Recovery will never be consumed on a released channel; unblock anyone gating on it. + stateConsumedFuture.complete(null); Review Comment: Why `complete(null)` rather than `completeExceptionally(...)` , similar to `upstreamReady`? ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannel.java: ########## @@ -123,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 = new CountDownLatch(1); Review Comment: Why is this a `CountDownLatch` here but a `CompletableFuture` for the equivalent state in `LocalInputChannel`? And more fundamentally, why do we need to block on upstream readiness at all here (as opposed to `LocalInputChannel`, which apparently doesn't need to)? ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannel.java: ########## @@ -188,6 +217,11 @@ void setExpectedSequenceNumber(int expectedSequenceNumber) { this.expectedSequenceNumber = expectedSequenceNumber; } + @VisibleForTesting + void completeUpstreamReadyForTest() { Review Comment: Could we avoid adding this test-only method by passing the latch (or future) itself out for tests to await on, instead? ditto `LocalInputChannel` ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RecoveredInputChannel.java: ########## @@ -115,23 +115,49 @@ public void setChannelStateWriter(ChannelStateWriter channelStateWriter) { this.channelStateWriter = checkNotNull(channelStateWriter); } - public final InputChannel toInputChannel() throws IOException { - Preconditions.checkState( - bufferFilteringCompleteFuture.isDone(), "buffer filtering is not complete"); - if (!inputGate.isCheckpointingDuringRecoveryEnabled()) { - Preconditions.checkState( - stateConsumedFuture.isDone(), "recovered state is not fully consumed"); + public final InputChannel toInputChannel(boolean needsRecovery) throws IOException { + if (needsRecovery) { + return toInputChannelInRecovery(); + } + synchronized (receivedBuffers) { + Preconditions.checkState(receivedBuffers.isEmpty(), "Received buffer should be empty."); } - // Extract remaining buffers before conversion. - // These buffers have been filtered but not yet consumed by the Task. + final InputChannel inputChannel = toInputChannelInternal(needsRecovery); + inputChannel.setup(); + inputChannel.checkpointStopped(lastStoppedCheckpointId); + return inputChannel; + } + + /** + * FLINK-38544 transitional: removed when the spilling backend lands. Creates the physical + * channel in recovery state and synchronously hands every queued recovered buffer over through + * the push interface. The legacy {@link EndOfInputChannelStateEvent} in the queue is dropped in + * translation; the {@link EndOfFetchedChannelStateEvent} sentinel takes its place. The sentinel + * is appended directly instead of via {@link + * RecoverableInputChannel#finishRecoveredBufferDelivery()} because that method waits for + * upstream readiness, which cannot happen while the mailbox thread is still converting channels + * (partitions are requested only after conversion). + */ + private InputChannel toInputChannelInRecovery() throws IOException { final ArrayDeque<Buffer> remainingBuffers; synchronized (receivedBuffers) { remainingBuffers = new ArrayDeque<>(receivedBuffers); receivedBuffers.clear(); } - final InputChannel inputChannel = toInputChannelInternal(remainingBuffers); + final InputChannel inputChannel = toInputChannelInternal(true); + inputChannel.setup(); + final RecoverableInputChannel recoverableChannel = (RecoverableInputChannel) inputChannel; + for (Buffer buffer : remainingBuffers) { + if (isEndOfInputChannelStateEvent(buffer)) { + buffer.recycleBuffer(); Review Comment: Don't we need a `break;` here once we've seen the `EndOfInputChannelStateEvent`? Or is it guaranteed that no further entries follow it in `remainingBuffers`? ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannel.java: ########## @@ -596,6 +754,11 @@ public void onBuffer(Buffer buffer, int sequenceNumber, int backlog, int subpart 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(); Review Comment: Did we run the the relevant benchmarks or the full suite with this change to confirm there's no regression? (on whatever PR version) (`countDown ` here is one `volatile` read IIUC) ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannel.java: ########## @@ -201,6 +235,119 @@ void setup() throws IOException { 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 Review Comment: nit: refomat comments? ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannel.java: ########## @@ -713,30 +889,110 @@ private void checkAnnouncedOnlyOnce(SequenceBuffer sequenceBuffer) { } /** - * 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); + } + } + + /** + * 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. + * + * @throws IOException 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<>(); + 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(); + if (isRecoveryCheckpointBarrier(sb.buffer, checkpointId)) { + sentinel = sb; + break; + } + // Skip non-data events (e.g. the EndOfFetchedChannelStateEvent sentinel appended + // after the recovered buffers): only recovered data buffers are snapshotted. + if (sb.buffer.isBuffer()) { Review Comment: Shouldn't we break or throw here if we encounter the `EndOfFetchedChannelStateEvent` sentinel before finding the matching `RecoveryCheckpointBarrier`? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
