From 8542b7853f80679466a1e5cd4406f3dacdb18533 Mon Sep 17 00:00:00 2001
From: Anton Kalashnikov <kaa.dev@yandex.ru>
Date: Tue, 7 Sep 2021 18:58:59 +0200
Subject: [PATCH] [refactor][runtime] Moved the INITIALIZING transition under
 cleanUp try-catch block

---
 .../flink/runtime/taskmanager/Task.java       | 24 +++++++++----------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java
index 76f00840672a89..1e0b1f7d1f10d2 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java
@@ -751,18 +751,6 @@ private void doRun() {
             // by the time we switched to running.
             this.invokable = invokable;
 
-            // switch to the INITIALIZING state, if that fails, we have been canceled/failed in the
-            // meantime
-            if (!transitionState(ExecutionState.DEPLOYING, ExecutionState.INITIALIZING)) {
-                throw new CancelTaskException();
-            }
-
-            taskManagerActions.updateTaskExecutionState(
-                    new TaskExecutionState(executionId, ExecutionState.INITIALIZING));
-
-            // make sure the user code classloader is accessible thread-locally
-            executingThread.setContextClassLoader(userCodeClassLoader.asClassLoader());
-
             restoreAndInvoke(invokable);
 
             // make sure, we enter the catch block if the task leaves the invoke() method due
@@ -924,6 +912,18 @@ else if (transitionState(current, ExecutionState.FAILED, t)) {
 
     private void restoreAndInvoke(TaskInvokable finalInvokable) throws Exception {
         try {
+            // switch to the INITIALIZING state, if that fails, we have been canceled/failed in the
+            // meantime
+            if (!transitionState(ExecutionState.DEPLOYING, ExecutionState.INITIALIZING)) {
+                throw new CancelTaskException();
+            }
+
+            taskManagerActions.updateTaskExecutionState(
+                    new TaskExecutionState(executionId, ExecutionState.INITIALIZING));
+
+            // make sure the user code classloader is accessible thread-locally
+            executingThread.setContextClassLoader(userCodeClassLoader.asClassLoader());
+
             runWithSystemExitMonitoring(finalInvokable::restore);
 
             if (!transitionState(ExecutionState.INITIALIZING, ExecutionState.RUNNING)) {