Update Worker.scala

apache · Jul 25, 2023 · 4d0d80b · 4d0d80b
1 parent 9139603
commit 4d0d80b
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
@@ -439,7 +439,6 @@ private[celeborn] class Worker(
       logInfo("Worker is stopped.")
       stopped = true
     }
-    shutdown.set(true)
   }
 
   private def registerWithMaster(): Unit = {
@@ -623,11 +622,15 @@ private[celeborn] class Worker(
   ShutdownHookManager.get().addShutdownHook(
     new Thread(new Runnable {
       override def run(): Unit = {
-        logInfo("Shutdown hook called.")
-        if (exitKind == CelebornExitKind.WORKER_GRACEFUL_SHUTDOWN) {
-          shutdownGracefully()
+        if (stopped) {
+          logInfo("Worker already stopped before call ShutdownHook.")
         } else {
-          exitImmediately()
+          logInfo("Shutdown hook called.")
+          if (exitKind == CelebornExitKind.WORKER_GRACEFUL_SHUTDOWN) {
+            shutdownGracefully()
+          } else {
+            exitImmediately()
+          }
         }
       }
     }),