Add tests.

JoshRosen · Jun 5, 2024 · 2c0b148 · 2c0b148
1 parent 34ac7de
commit 2c0b148
Show file tree

Hide file tree

Showing 2 changed files with 72 additions and 0 deletions.
diff --git a/core/src/main/scala/org/apache/spark/internal/config/package.scala b/core/src/main/scala/org/apache/spark/internal/config/package.scala
@@ -2206,6 +2206,12 @@ package object config {
     .toSequence
     .createWithDefault(Nil)
 
+  private[spark] val SUBMIT_CALL_SYSTEM_EXIT_ON_MAIN_EXIT =
+    ConfigBuilder("spark.submit.callSystemExitOnMainExit")
+      .version("4.0.0")
+      .booleanConf
+      .createWithDefault(false)
+
   private[spark] val SCHEDULER_ALLOCATION_FILE =
     ConfigBuilder("spark.scheduler.allocation.file")
       .version("0.8.1")

diff --git a/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala b/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
@@ -1589,6 +1589,44 @@ class SparkSubmitSuite
       runSparkSubmit(argsSuccess, expectFailure = false))
   }
 
+  test("spark.submit.callSystemExitOnMainExit returns non-zero exit code on unclean main exit") {
+    val unusedJar = TestUtils.createJarWithClasses(Seq.empty)
+    val args = Seq(
+      "--class", MainThrowsUncaughtExceptionSparkApplicationTest.getClass.getName.stripSuffix("$"),
+      "--name", "testApp",
+      "--conf", s"${SUBMIT_CALL_SYSTEM_EXIT_ON_MAIN_EXIT.key}=true",
+      unusedJar.toString
+    )
+    assertResult(1)(runSparkSubmit(args, expectFailure = true))
+  }
+
+  test("spark.submit.callSystemExitOnMainExit calls system exit on clean main exit") {
+    val unusedJar = TestUtils.createJarWithClasses(Seq.empty)
+    val args = Seq(
+      "--class", NonDaemonThreadSparkApplicationTest.getClass.getName.stripSuffix("$"),
+      "--name", "testApp",
+      "--conf", s"${SUBMIT_CALL_SYSTEM_EXIT_ON_MAIN_EXIT.key}=true",
+      unusedJar.toString
+    )
+    // With SUBMIT_CALL_SYSTEM_EXIT_ON_MAIN_EXIT set to false, the non-daemon thread will
+    // prevent the JVM from beginning shutdown and the following call will fail with a
+    // timeout:
+    assertResult(0)(runSparkSubmit(args))
+  }
+
+  test("spark.submit.callSystemExitOnMainExit with main that explicitly calls System.exit") {
+    val unusedJar = TestUtils.createJarWithClasses(Seq.empty)
+    val args = Seq(
+      "--class",
+      MainExplicitlyCallsSystemExit3SparkApplicationTest.getClass.getName.stripSuffix("$"),
+      "--name", "testApp",
+      "--conf", s"${SUBMIT_CALL_SYSTEM_EXIT_ON_MAIN_EXIT.key}=true",
+      unusedJar.toString
+    )
+    // This main class explicitly exits with System.exit(3), hence this expected exit code:
+    assertResult(3)(runSparkSubmit(args, expectFailure = true))
+  }
+
   private def testRemoteResources(
       enableHttpFs: Boolean,
       forceDownloadSchemes: Seq[String] = Nil): Unit = {
@@ -1855,6 +1893,34 @@ object SimpleApplicationTest {
   }
 }
 
+object MainThrowsUncaughtExceptionSparkApplicationTest {
+  def main(args: Array[String]): Unit = {
+    throw new Exception("User exception")
+  }
+}
+
+object NonDaemonThreadSparkApplicationTest {
+  def main(args: Array[String]): Unit = {
+    val nonDaemonThread: Thread = new Thread {
+      override def run(): Unit = {
+        while (true) {
+          Thread.sleep(1000)
+        }
+      }
+    }
+    nonDaemonThread.setDaemon(false)
+    nonDaemonThread.setName("Non-Daemon-Thread")
+    nonDaemonThread.start()
+    // Fall off the end of the main method.
+  }
+}
+
+object MainExplicitlyCallsSystemExit3SparkApplicationTest {
+  def main(args: Array[String]): Unit = {
+    System.exit(3)
+  }
+}
+
 object UserClasspathFirstTest {
   def main(args: Array[String]): Unit = {
     val ccl = Thread.currentThread().getContextClassLoader()