[FLINK-38543][checkpoint] Fix Mailbox loop interrupted before recovery finished

Author: 1996fanrui

flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java

@@ -913,8 +913,14 @@ private CompletableFuture<Void> restoreStateAndGates(
                                             "Input gate request partitions"));
         }
 
-        return CompletableFuture.allOf(recoveredFutures.toArray(new CompletableFuture[0]))
-                .thenRun(mailboxProcessor::suspend);
+        // Return allOf result instead of thenRun result.
+        // thenRun returns a NEW future that completes after the callback finishes.
+        // Since suspend() runs on the async thread and just sends a poison mail,
+        // the mailbox loop can exit before suspend() returns, causing isDone() to be false.
+        CompletableFuture<Void> allRecoveredFuture =
+                CompletableFuture.allOf(recoveredFutures.toArray(new CompletableFuture[0]));
+        allRecoveredFuture.thenRun(mailboxProcessor::suspend);
+        return allRecoveredFuture;
     }
 
     private void ensureNotCanceled() {