1996fanrui commented on code in PR #28662:
URL: https://github.com/apache/flink/pull/28662#discussion_r3777646067


##########
flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateDrainer.java:
##########
@@ -18,71 +18,203 @@
 package org.apache.flink.runtime.checkpoint.channel;
 
 import org.apache.flink.annotation.Internal;
+import 
org.apache.flink.runtime.checkpoint.channel.FetchedChannelStateReader.SpillSegment;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
 import 
org.apache.flink.runtime.io.network.partition.consumer.RecoverableInputChannel;
 
 import java.io.Closeable;
 import java.io.IOException;
+import java.io.InputStream;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 
+import static org.apache.flink.util.Preconditions.checkArgument;
 import static org.apache.flink.util.Preconditions.checkNotNull;
 
 /**
- * Drains a {@link FetchedChannelState} into the physical recovered channels 
and inserts {@link
- * RecoveryCheckpointBarrier}s when a checkpoint fires during recovery.
+ * Drains a {@link FetchedChannelState} into recovered-buffer queues and 
snapshots remaining
+ * segments when a checkpoint fires during recovery.
  *
- * <p>FLINK-38544 transitional in-memory implementation: the in-memory 
recovery backend already
- * pushed every recovered buffer into the physical channels' own queues at 
conversion time (via
- * {@code requestPartitions(true)}), so draining is just appending the 
end-of-recovered-state
- * sentinel to each channel, and there is never an undrained residue to 
snapshot — inserting the
- * barrier into the in-recovery channels is enough. The disk-based drainer of 
the spilling backend
- * replaces this, reading segments off spill files and returning a reader over 
the undrained slice.
+ * <p>The drainer lock pairs channel delivery with reader-cursor advancement 
and also protects
+ * snapshot creation plus barrier insertion. Disk reads and buffer allocation 
stay outside that
+ * lock.
  */
 @Internal
 public final class FetchedChannelStateDrainer implements 
RecoveryCheckpointTrigger, Closeable {
 
+    private final FetchedChannelStateReader rootReader;
+
+    private final ResolvedChannels channels;
+
+    private final Object lock = new Object();
     private final FetchedChannelState channelState;
 
-    private final List<RecoverableInputChannel> channels;
+    /**
+     * Set under {@link #lock} once {@link #drain()} has consumed every 
segment. After that the
+     * {@link #rootReader} is closed by {@link #close()}, so a later {@link
+     * #snapshotAndInsertBarriers} must not derive from it; it returns an 
empty reader instead.
+     * Guarded by the lock so the check is atomic with barrier insertion.
+     */
+    private boolean drainFinished;
 
     public FetchedChannelStateDrainer(
             FetchedChannelState channelState, List<RecoverableInputChannel> 
channels) {
-        this.channelState = checkNotNull(channelState);
-        this.channels = checkNotNull(channels);
+        this.channelState = channelState;
+        this.rootReader = checkNotNull(channelState).reader();
+        this.channels = new ResolvedChannels(channels);
+    }
+
+    private static final class ResolvedChannels {
+        final List<RecoverableInputChannel> allChannels;
+        final Map<InputChannelInfo, RecoverableInputChannel> channelByInfo;
+
+        ResolvedChannels(List<RecoverableInputChannel> all) {
+            this.allChannels = all;
+            Map<InputChannelInfo, RecoverableInputChannel> byInfo = new 
HashMap<>();
+            for (RecoverableInputChannel ch : all) {
+                byInfo.put(ch.getChannelInfo(), ch);
+            }
+            this.channelByInfo = byInfo;
+        }
     }
 
     /**
-     * Appends the end-of-recovered-state sentinel to every converted channel. 
Each channel first
-     * waits for its upstream to be ready, so this must run on the 
channelIOExecutor rather than
-     * block the mailbox thread. Only once the sentinel is in place can the 
consume path flip the
-     * channel out of recovery, which guarantees live data is never polled 
before the upstream
-     * connection exists.
+     * Drains all segments from the spill file into the corresponding recovery 
buffer queues. Each
+     * segment is split into chunks of at most {@code memorySegmentSize} 
bytes; a full chunk is
+     * delivered under the drainer lock paired with a segment commit. After 
all segments are
+     * drained, every channel's {@link 
RecoverableInputChannel#finishRecoveredBufferDelivery()} is
+     * called.
+     *
+     * <p>Disk reads and buffer allocations happen outside the lock; only the 
"deliver + commit"
+     * pair is locked to guarantee atomicity with snapshot.
      */
     public void drain() throws IOException, InterruptedException {
         channelState.release();
-        for (RecoverableInputChannel channel : channels) {
-            channel.finishRecoveredBufferDelivery();
+        Optional<SpillSegment> next;
+        while ((next = rootReader.advanceAndGetNextSegment()).isPresent()) {
+            SpillSegment seg = next.get();
+            RecoverableInputChannel ch = 
channels.channelByInfo.get(seg.channelInfo());
+            if (ch == null) {
+                throw new IllegalStateException(
+                        "Drain: no physical channel found for " + 
seg.channelInfo());
+            }
+            drainSegment(seg, ch);
+        }
+
+        // Mark drain done before rootReader is closed, so a concurrent 
snapshot returns empty
+        // rather than deriving from the soon-to-be-closed rootReader. Under 
the lock to stay atomic
+        // with snapshotAndInsertBarriers' check.
+        synchronized (lock) {
+            drainFinished = true;
+        }
+        for (RecoverableInputChannel ch : channels.allChannels) {
+            ch.finishRecoveredBufferDelivery();
         }
     }
 
     /**
-     * Inserts a {@link RecoveryCheckpointBarrier} into every channel that is 
still in recovery, so
-     * that {@code checkpointStarted}'s in-recovery branch can persist exactly 
the pre-barrier
-     * recovered data. There is no snapshot side for the in-memory backend: 
everything a checkpoint
-     * must persist is already inside the channels' queues, so the snapshot is 
inherently empty.
+     * Drains one segment into the given channel. Fills buffers from the 
segment's opaque byte
+     * stream in chunks of at most {@code memorySegmentSize} bytes. A full 
buffer is delivered under
+     * the lock and a fresh one is requested; a partial tail buffer (if 
non-empty) is also
+     * delivered.
+     */
+    private void drainSegment(SpillSegment seg, RecoverableInputChannel ch)
+            throws IOException, InterruptedException {
+        InputStream in = seg.bodyStream();
+        int remaining = seg.length();
+        Buffer buf = ch.requestRecoveryBufferBlocking();
+        try {
+            int cap = buf.getMaxCapacity();
+
+            while (fill(buf, in, cap - buf.getSize()) > 0) {
+                if (buf.getSize() == cap) {
+                    // Buffer is full: deliver under lock and request a fresh 
one.
+                    remaining -= cap;
+                    Buffer full = buf;
+                    buf = null;
+                    synchronized (lock) {
+                        ch.onRecoveredStateBuffer(full);
+                        seg.commit();
+                    }
+                    if (remaining == 0) {
+                        // The segment ends on this buffer boundary: no point 
in requesting a
+                        // buffer that the EOF below would immediately recycle.
+                        return;
+                    }
+                    buf = ch.requestRecoveryBufferBlocking();
+                    cap = buf.getMaxCapacity();
+                }
+                // If buf is not full yet, the fill returned > 0 bytes but 
segment is not exhausted;
+                // loop and keep filling the same buffer.
+            }
+
+            if (buf.getSize() > 0) {
+                // Deliver the partial tail buffer.
+                Buffer tail = buf;
+                buf = null;
+                synchronized (lock) {
+                    ch.onRecoveredStateBuffer(tail);
+                    seg.commit();
+                }

Review Comment:
   Applied, plus the tail delivery after the loop (otherwise the last partial 
buffer is lost) and the try/catch for the in-flight buffer.



##########
flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/FetchedChannelStateDrainer.java:
##########
@@ -18,71 +18,203 @@
 package org.apache.flink.runtime.checkpoint.channel;
 
 import org.apache.flink.annotation.Internal;
+import 
org.apache.flink.runtime.checkpoint.channel.FetchedChannelStateReader.SpillSegment;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
 import 
org.apache.flink.runtime.io.network.partition.consumer.RecoverableInputChannel;
 
 import java.io.Closeable;
 import java.io.IOException;
+import java.io.InputStream;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 
+import static org.apache.flink.util.Preconditions.checkArgument;
 import static org.apache.flink.util.Preconditions.checkNotNull;
 
 /**
- * Drains a {@link FetchedChannelState} into the physical recovered channels 
and inserts {@link
- * RecoveryCheckpointBarrier}s when a checkpoint fires during recovery.
+ * Drains a {@link FetchedChannelState} into recovered-buffer queues and 
snapshots remaining
+ * segments when a checkpoint fires during recovery.
  *
- * <p>FLINK-38544 transitional in-memory implementation: the in-memory 
recovery backend already
- * pushed every recovered buffer into the physical channels' own queues at 
conversion time (via
- * {@code requestPartitions(true)}), so draining is just appending the 
end-of-recovered-state
- * sentinel to each channel, and there is never an undrained residue to 
snapshot — inserting the
- * barrier into the in-recovery channels is enough. The disk-based drainer of 
the spilling backend
- * replaces this, reading segments off spill files and returning a reader over 
the undrained slice.
+ * <p>The drainer lock pairs channel delivery with reader-cursor advancement 
and also protects
+ * snapshot creation plus barrier insertion. Disk reads and buffer allocation 
stay outside that
+ * lock.
  */
 @Internal
 public final class FetchedChannelStateDrainer implements 
RecoveryCheckpointTrigger, Closeable {
 
+    private final FetchedChannelStateReader rootReader;
+
+    private final ResolvedChannels channels;
+
+    private final Object lock = new Object();
     private final FetchedChannelState channelState;
 
-    private final List<RecoverableInputChannel> channels;
+    /**
+     * Set under {@link #lock} once {@link #drain()} has consumed every 
segment. After that the
+     * {@link #rootReader} is closed by {@link #close()}, so a later {@link
+     * #snapshotAndInsertBarriers} must not derive from it; it returns an 
empty reader instead.
+     * Guarded by the lock so the check is atomic with barrier insertion.
+     */
+    private boolean drainFinished;
 
     public FetchedChannelStateDrainer(
             FetchedChannelState channelState, List<RecoverableInputChannel> 
channels) {
-        this.channelState = checkNotNull(channelState);
-        this.channels = checkNotNull(channels);
+        this.channelState = channelState;
+        this.rootReader = checkNotNull(channelState).reader();
+        this.channels = new ResolvedChannels(channels);
+    }
+
+    private static final class ResolvedChannels {
+        final List<RecoverableInputChannel> allChannels;
+        final Map<InputChannelInfo, RecoverableInputChannel> channelByInfo;
+
+        ResolvedChannels(List<RecoverableInputChannel> all) {
+            this.allChannels = all;
+            Map<InputChannelInfo, RecoverableInputChannel> byInfo = new 
HashMap<>();
+            for (RecoverableInputChannel ch : all) {
+                byInfo.put(ch.getChannelInfo(), ch);
+            }
+            this.channelByInfo = byInfo;
+        }
     }
 
     /**
-     * Appends the end-of-recovered-state sentinel to every converted channel. 
Each channel first
-     * waits for its upstream to be ready, so this must run on the 
channelIOExecutor rather than
-     * block the mailbox thread. Only once the sentinel is in place can the 
consume path flip the
-     * channel out of recovery, which guarantees live data is never polled 
before the upstream
-     * connection exists.
+     * Drains all segments from the spill file into the corresponding recovery 
buffer queues. Each
+     * segment is split into chunks of at most {@code memorySegmentSize} 
bytes; a full chunk is
+     * delivered under the drainer lock paired with a segment commit. After 
all segments are
+     * drained, every channel's {@link 
RecoverableInputChannel#finishRecoveredBufferDelivery()} is
+     * called.
+     *
+     * <p>Disk reads and buffer allocations happen outside the lock; only the 
"deliver + commit"
+     * pair is locked to guarantee atomicity with snapshot.
      */
     public void drain() throws IOException, InterruptedException {
         channelState.release();
-        for (RecoverableInputChannel channel : channels) {
-            channel.finishRecoveredBufferDelivery();
+        Optional<SpillSegment> next;
+        while ((next = rootReader.advanceAndGetNextSegment()).isPresent()) {
+            SpillSegment seg = next.get();
+            RecoverableInputChannel ch = 
channels.channelByInfo.get(seg.channelInfo());
+            if (ch == null) {
+                throw new IllegalStateException(
+                        "Drain: no physical channel found for " + 
seg.channelInfo());
+            }
+            drainSegment(seg, ch);
+        }
+
+        // Mark drain done before rootReader is closed, so a concurrent 
snapshot returns empty
+        // rather than deriving from the soon-to-be-closed rootReader. Under 
the lock to stay atomic
+        // with snapshotAndInsertBarriers' check.
+        synchronized (lock) {
+            drainFinished = true;
+        }
+        for (RecoverableInputChannel ch : channels.allChannels) {
+            ch.finishRecoveredBufferDelivery();
         }
     }
 
     /**
-     * Inserts a {@link RecoveryCheckpointBarrier} into every channel that is 
still in recovery, so
-     * that {@code checkpointStarted}'s in-recovery branch can persist exactly 
the pre-barrier
-     * recovered data. There is no snapshot side for the in-memory backend: 
everything a checkpoint
-     * must persist is already inside the channels' queues, so the snapshot is 
inherently empty.
+     * Drains one segment into the given channel. Fills buffers from the 
segment's opaque byte
+     * stream in chunks of at most {@code memorySegmentSize} bytes. A full 
buffer is delivered under
+     * the lock and a fresh one is requested; a partial tail buffer (if 
non-empty) is also
+     * delivered.
+     */
+    private void drainSegment(SpillSegment seg, RecoverableInputChannel ch)
+            throws IOException, InterruptedException {
+        InputStream in = seg.bodyStream();
+        int remaining = seg.length();
+        Buffer buf = ch.requestRecoveryBufferBlocking();
+        try {
+            int cap = buf.getMaxCapacity();
+
+            while (fill(buf, in, cap - buf.getSize()) > 0) {
+                if (buf.getSize() == cap) {
+                    // Buffer is full: deliver under lock and request a fresh 
one.
+                    remaining -= cap;
+                    Buffer full = buf;
+                    buf = null;
+                    synchronized (lock) {
+                        ch.onRecoveredStateBuffer(full);
+                        seg.commit();
+                    }
+                    if (remaining == 0) {
+                        // The segment ends on this buffer boundary: no point 
in requesting a
+                        // buffer that the EOF below would immediately recycle.
+                        return;
+                    }
+                    buf = ch.requestRecoveryBufferBlocking();
+                    cap = buf.getMaxCapacity();
+                }
+                // If buf is not full yet, the fill returned > 0 bytes but 
segment is not exhausted;
+                // loop and keep filling the same buffer.
+            }
+
+            if (buf.getSize() > 0) {
+                // Deliver the partial tail buffer.
+                Buffer tail = buf;
+                buf = null;
+                synchronized (lock) {
+                    ch.onRecoveredStateBuffer(tail);
+                    seg.commit();
+                }
+            } else {
+                buf.recycleBuffer();
+                buf = null;

Review Comment:
   It can't, and it's gone — the rewrite below removed that branch entirely.



##########
flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateCheckpointWriter.java:
##########
@@ -161,6 +163,52 @@ void writeInput(
         }
     }
 
+    void writeInputFromSpill(
+            JobVertexID jobVertexID, int subtaskIndex, 
FetchedChannelStateReader reader) {
+        if (isDone()) {
+            try {
+                reader.close();
+            } catch (Exception e) {
+                LOG.info(
+                        "Failed to close the fetched channel state reader of 
checkpoint {}",
+                        checkpointId,
+                        e);
+            }
+            return;
+        }
+        ChannelStatePendingResult pendingResult =
+                getChannelStatePendingResult(jobVertexID, subtaskIndex);
+        runWithChecks(
+                () -> {
+                    checkState(!pendingResult.isAllInputsReceived());

Review Comment:
   Right. Moved the close into a method-level `finally`, which also covers the 
two earlier throw sites (`getChannelStatePendingResult` and `runWithChecks`' 
`checkState`).



-- 
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]

Reply via email to