[flink] 02/02: [FLINK-29816][streaming] Fix the bug that StreamTask doesn't handle exception during restoring

[gu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

guoweijie pushed a commit to branch release-1.17
in repository https://gitbox.apache.org/repos/asf/flink.git

commit 65e455189d2a88fd62be2e8a464feab4687ee088
Author: Roc Marshal <fl...@126.com>
AuthorDate: Sun Feb 26 12:18:00 2023 +0800

    [FLINK-29816][streaming] Fix the bug that StreamTask doesn't handle exception during restoring
    
    (cherry picked from commit a98bb9a6c978d5458cf4cc11dc7b68034f4def4b)
---
 .../flink/streaming/runtime/tasks/StreamTask.java  |  9 ++++--
 .../streaming/runtime/tasks/StreamTaskTest.java    | 37 ++++++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
index 6b52ed5da3b..d2d903fc33a 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
@@ -258,6 +258,9 @@ public abstract class StreamTask<OUT, OP extends StreamOperator<OUT>>
      */
     private volatile boolean isRunning;
 
+    /** Flag to mark the task at restoring duration in {@link #restore()}. */
+    private volatile boolean isRestoring;
+
     /** Flag to mark this task as canceled. */
     private volatile boolean canceled;
 
@@ -677,6 +680,7 @@ public abstract class StreamTask<OUT, OP extends StreamOperator<OUT>>
             LOG.debug("Re-restore attempt rejected.");
             return;
         }
+        isRestoring = true;
         closedOperators = false;
         LOG.debug("Initializing {}.", getName());
 
@@ -718,6 +722,7 @@ public abstract class StreamTask<OUT, OP extends StreamOperator<OUT>>
         channelIOExecutor.shutdown();
 
         isRunning = true;
+        isRestoring = false;
     }
 
     private CompletableFuture<Void> restoreGates() throws Exception {
@@ -1551,8 +1556,8 @@ public abstract class StreamTask<OUT, OP extends StreamOperator<OUT>>
      */
     @Override
     public void handleAsyncException(String message, Throwable exception) {
-        if (isRunning) {
-            // only fail if the task is still running
+        if (isRestoring || isRunning) {
+            // only fail if the task is still in restoring or running
             asyncExceptionHandler.handleAsyncException(message, exception);
         }
     }
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java
index 51724eb3113..585cd51ee17 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java
@@ -402,6 +402,43 @@ public class StreamTaskTest {
         }
     }
 
+    @Test
+    void testHandleAsyncExceptionDuringRestoring() throws Exception {
+        MockEnvironment mockEnvironment = MockEnvironment.builder().build();
+        Throwable expectedException = new RuntimeException("RUNTIME EXCEPTION");
+
+        mockEnvironment.setExpectedExternalFailureCause(AsynchronousException.class);
+        final String expectedErrorMessage = "EXPECTED_ERROR MESSAGE";
+
+        StreamTaskITCase.NoOpStreamTask initThrowExceptionTask =
+                new StreamTaskITCase.NoOpStreamTask(mockEnvironment) {
+
+                    @Override
+                    protected void init() throws Exception {
+                        super.init();
+
+                        // Throw exception during restoring.
+                        CompletableFuture.runAsync(
+                                        () ->
+                                                this.handleAsyncException(
+                                                        expectedErrorMessage, expectedException))
+                                .get();
+                    }
+                };
+        initThrowExceptionTask.restore();
+
+        Optional<? extends Throwable> actualExternalFailureCause =
+                mockEnvironment.getActualExternalFailureCause();
+        final Throwable actualException =
+                actualExternalFailureCause.orElseThrow(
+                        () -> new AssertionError("Expected exceptional completion"));
+
+        assertThat(actualException)
+                .isInstanceOf(AsynchronousException.class)
+                .hasMessage(expectedErrorMessage)
+                .hasCause(expectedException);
+    }
+
     /**
      * This test checks the async exceptions handling wraps the message and cause as an
      * AsynchronousException and propagates this to the environment.