[FLINK-38543][checkpoint] Introduce bufferFilteringCompleteFuture for earlier RUNNING state transition

Author: 1996fanrui

flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/InputGate.java

@@ -192,5 +192,12 @@ public String toString() {
 
     public abstract CompletableFuture<Void> getStateConsumedFuture();
 
+    /**
+     * Returns a future that completes when buffer filtering is complete for all channels. This
+     * future completes before {@link #getStateConsumedFuture()}, enabling earlier RUNNING state
+     * transition when unaligned checkpoint during recovery is enabled.
+     */
+    public abstract CompletableFuture<Void> getBufferFilteringCompleteFuture();
+
     public abstract void finishReadRecoveredState() throws IOException;
 }

flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RecoveredInputChannel.java

@@ -62,6 +62,13 @@ public abstract class RecoveredInputChannel extends InputChannel implements Chan
     private final CompletableFuture<?> stateConsumedFuture = new CompletableFuture<>();
     protected final BufferManager bufferManager;
 
+    /**
+     * Future that completes when recovered buffers have been filtered for this channel. This
+     * completes before stateConsumedFuture, enabling earlier RUNNING state transition when
+     * unaligned checkpoint during recovery is enabled.
+     */
+    private final CompletableFuture<Void> bufferFilteringCompleteFuture = new CompletableFuture<>();
+
     @GuardedBy("receivedBuffers")
     private boolean isReleased;
 
@@ -109,8 +116,16 @@ public void setChannelStateWriter(ChannelStateWriter channelStateWriter) {
     }
 
     public final InputChannel toInputChannel() throws IOException {
-        Preconditions.checkState(
-                stateConsumedFuture.isDone(), "recovered state is not fully consumed");
+        // Check the appropriate future based on configuration:
+        // - When unaligned during recovery is enabled: check bufferFilteringCompleteFuture
+        // - When disabled: check stateConsumedFuture (original behavior)
+        if (inputGate.isCheckpointingDuringRecoveryEnabled()) {
+            Preconditions.checkState(
+                    bufferFilteringCompleteFuture.isDone(), "buffer filtering is not complete");
+        } else {
+            Preconditions.checkState(
+                    stateConsumedFuture.isDone(), "recovered state is not fully consumed");
+        }
 
         // Extract remaining buffers before conversion.
         // These buffers have been filtered but not yet consumed by the Task.
@@ -140,6 +155,14 @@ public void checkpointStopped(long checkpointId) {
     protected abstract InputChannel toInputChannelInternal(ArrayDeque<Buffer> remainingBuffers)
             throws IOException;
 
+    /**
+     * Returns the future that completes when buffer filtering is complete. This future completes
+     * before stateConsumedFuture, at the point when finishReadRecoveredState() is called.
+     */
+    CompletableFuture<Void> getBufferFilteringCompleteFuture() {
+        return bufferFilteringCompleteFuture;
+    }
+
     CompletableFuture<?> getStateConsumedFuture() {
         return stateConsumedFuture;
     }
@@ -180,6 +203,13 @@ public void finishReadRecoveredState() throws IOException {
                 EventSerializer.toBuffer(EndOfInputChannelStateEvent.INSTANCE, false));
         bufferManager.releaseFloatingBuffers();
         LOG.debug("{}/{} finished recovering input.", inputGate.getOwningTaskName(), channelInfo);
+
+        // Complete bufferFilteringCompleteFuture only when unaligned during recovery is enabled.
+        // This signals that buffer filtering is complete, allowing earlier RUNNING state
+        // transition. When the config is disabled, this future should not be completed.
+        if (inputGate.isCheckpointingDuringRecoveryEnabled()) {
+            bufferFilteringCompleteFuture.complete(null);
+        }
     }
 
     @Nullable

flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/SingleInputGate.java

@@ -341,6 +341,21 @@ public boolean isCheckpointingDuringRecoveryEnabled() {
         return checkpointingDuringRecoveryEnabled;
     }
 
+    @Override
+    public CompletableFuture<Void> getBufferFilteringCompleteFuture() {
+        synchronized (requestLock) {
+            List<CompletableFuture<?>> futures = new ArrayList<>(numberOfInputChannels);
+            for (InputChannel inputChannel : inputChannels()) {
+                if (inputChannel instanceof RecoveredInputChannel) {
+                    futures.add(
+                            ((RecoveredInputChannel) inputChannel)
+                                    .getBufferFilteringCompleteFuture());
+                }
+            }
+            return CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]));
+        }
+    }
+
     @Override
     public void requestPartitions() {
         synchronized (requestLock) {

flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/UnionInputGate.java

@@ -350,6 +350,15 @@ public CompletableFuture<Void> getStateConsumedFuture() {
                         .toArray(new CompletableFuture[] {}));
     }
 
+    @Override
+    public CompletableFuture<Void> getBufferFilteringCompleteFuture() {
+        return CompletableFuture.allOf(
+                inputGatesByGateIndex.values().stream()
+                        .map(InputGate::getBufferFilteringCompleteFuture)
+                        .collect(Collectors.toList())
+                        .toArray(new CompletableFuture[] {}));
+    }
+
     @Override
     public void requestPartitions() throws IOException {
         for (InputGate inputGate : inputGatesByGateIndex.values()) {

flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/InputGateWithMetrics.java

@@ -120,6 +120,11 @@ public CompletableFuture<Void> getStateConsumedFuture() {
         return inputGate.getStateConsumedFuture();
     }
 
+    @Override
+    public CompletableFuture<Void> getBufferFilteringCompleteFuture() {
+        return inputGate.getBufferFilteringCompleteFuture();
+    }
+
     @Override
     public void requestPartitions() throws IOException {
         inputGate.requestPartitions();

flink-runtime/src/test/java/org/apache/flink/streaming/runtime/io/MockIndexedInputGate.java

@@ -57,6 +57,11 @@ public CompletableFuture<Void> getStateConsumedFuture() {
         return CompletableFuture.completedFuture(null);
     }
 
+    @Override
+    public CompletableFuture<Void> getBufferFilteringCompleteFuture() {
+        return CompletableFuture.completedFuture(null);
+    }
+
     @Override
     public void finishReadRecoveredState() {}
 

flink-runtime/src/test/java/org/apache/flink/streaming/runtime/io/MockInputGate.java

@@ -80,6 +80,11 @@ public CompletableFuture<Void> getStateConsumedFuture() {
         return CompletableFuture.completedFuture(null);
     }
 
+    @Override
+    public CompletableFuture<Void> getBufferFilteringCompleteFuture() {
+        return CompletableFuture.completedFuture(null);
+    }
+
     @Override
     public void finishReadRecoveredState() {}
 

flink-runtime/src/test/java/org/apache/flink/streaming/runtime/io/checkpointing/AlignedCheckpointsMassiveRandomTest.java

@@ -263,6 +263,11 @@ public CompletableFuture<Void> getStateConsumedFuture() {
             return CompletableFuture.completedFuture(null);
         }
 
+        @Override
+        public CompletableFuture<Void> getBufferFilteringCompleteFuture() {
+            return CompletableFuture.completedFuture(null);
+        }
+
         @Override
         public void finishReadRecoveredState() {}