Pull detection of being in a shutdown hook into utility function.

2 files changed, 23 insertions, 14 deletions

diff --git a/core/src/main/scala/spark/Utils.scala b/core/src/main/scala/spark/Utils.scala
index 28d643abca..81daacf958 100644
--- a/core/src/main/scala/spark/Utils.scala
+++ b/core/src/main/scala/spark/Utils.scala
@@ -454,4 +454,25 @@ private object Utils extends Logging {
   def clone[T](value: T, serializer: SerializerInstance): T = {
     serializer.deserialize[T](serializer.serialize(value))
   }
+
+  /**
+   * Detect whether this thread might be executing a shutdown hook. Will always return true if
+   * the current thread is a running a shutdown hook but may spuriously return true otherwise (e.g.
+   * if System.exit was just called by a concurrent thread).
+   *
+   * Currently, this detects whether the JVM is shutting down by Runtime#addShutdownHook throwing
+   * an IllegalStateException.
+   */
+  def inShutdown(): Boolean = {
+    try {
+      val hook = new Thread {
+        override def run() {}
+      }
+      Runtime.getRuntime.addShutdownHook(hook)
+      Runtime.getRuntime.removeShutdownHook(hook)
+    } catch {
+      case ise: IllegalStateException => return true
+    }
+    return false
+  }
 }
diff --git a/core/src/main/scala/spark/executor/Executor.scala b/core/src/main/scala/spark/executor/Executor.scala
index b63bec11ad..5de09030aa 100644
--- a/core/src/main/scala/spark/executor/Executor.scala
+++ b/core/src/main/scala/spark/executor/Executor.scala
@@ -52,20 +52,8 @@ private[spark] class Executor extends Logging {
             logError("Uncaught exception in thread " + thread, exception)
             
             // We may have been called from a shutdown hook. If so, we must not call System.exit().
-            // (If we do, we will deadlock.) Runtime#addShutdownHook should fail if we are shutting
-            // down, which would either occur if we were called from a shutdown hook or if
-            // a System.exit() occured concurrently.
-            var shuttingDown = false
-            try {
-              val hook = new Thread {
-                override def run() {}
-              }
-              Runtime.getRuntime.addShutdownHook(hook)
-              Runtime.getRuntime.removeShutdownHook(hook)
-            } catch {
-              case ise: IllegalStateException => shuttingDown = true
-            }
-            if (!shuttingDown) {
+            // (If we do, we will deadlock.)
+            if (!Utils.inShutdown()) {
               if (exception.isInstanceOf[OutOfMemoryError]) {
                 System.exit(ExecutorExitCode.OOM)
               } else {