[KYUUBI #7590][SPARK] Exit engine when SparkContext stops

externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/SparkSQLEngine.scala

@@ -230,6 +230,8 @@ object SparkSQLEngine extends Logging {
 
   private val sparkSessionCreated = new AtomicBoolean(false)
 
+  private val ENGINE_TERMINATION_CHECK_INTERVAL_SECONDS = 1L
+
   final private val POD_NAME_MAX_LENGTH = 253
   final private val POD_UID_MAX_LENGTH = 36
   final private val POD_LOGS_DIRECTORY_SEPARATOR_LENGTH = 2
@@ -414,7 +416,11 @@ object SparkSQLEngine extends Logging {
         try {
           startEngine(spark)
           // blocking main thread
-          countDownLatch.await()
+          waitForEngineTermination(
+            countDownLatch,
+            () => spark.sparkContext.isStopped,
+            ENGINE_TERMINATION_CHECK_INTERVAL_SECONDS,
+            TimeUnit.SECONDS)
         } catch {
           case e: KyuubiException =>
             currentEngine match {
@@ -449,6 +455,21 @@ object SparkSQLEngine extends Logging {
     }
   }
 
+  @VisibleForTesting
+  private[kyuubi] def waitForEngineTermination(
+      latch: CountDownLatch,
+      isSparkContextStopped: () => Boolean,
+      checkInterval: Long,
+      timeUnit: TimeUnit): Unit = {
+    while (!latch.await(checkInterval, timeUnit)) {
+      // Engine shutdown may stall before stopServer releases the latch.
+      if (isSparkContextStopped() && !latch.await(0, TimeUnit.NANOSECONDS)) {
+        throw new KyuubiException(
+          "SparkContext stopped while engine is running, exiting main thread")
+      }
+    }
+  }
+
   private def startInitTimeoutChecker(startTime: Long, timeout: Long): Unit = {
     val mainThread = Thread.currentThread()
     val checker = new Thread(

externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/SparkSQLEngineSuite.scala

@@ -17,13 +17,42 @@
 
 package org.apache.kyuubi.engine.spark
 
-import org.apache.kyuubi.KyuubiFunSuite
+import java.util.concurrent.{CountDownLatch, TimeUnit}
+
+import org.apache.kyuubi.{KyuubiException, KyuubiFunSuite}
 
 class SparkSQLEngineSuite extends KyuubiFunSuite {
 
   private val namespace = "n" * 63
   private val podUid = "u" * 36
 
+  test("[KYUUBI #7590] stop waiting when SparkContext is stopped") {
+    val latch = new CountDownLatch(1)
+
+    val error = intercept[KyuubiException] {
+      SparkSQLEngine.waitForEngineTermination(
+        latch,
+        () => true,
+        1,
+        TimeUnit.MILLISECONDS)
+    }
+
+    assert(error.getMessage.contains("SparkContext stopped"))
+  }
+
+  test("[KYUUBI #7590] complete normally when shutdown wins the race") {
+    val latch = new CountDownLatch(1)
+
+    SparkSQLEngine.waitForEngineTermination(
+      latch,
+      () => {
+        latch.countDown()
+        true
+      },
+      1,
+      TimeUnit.MILLISECONDS)
+  }
+
   test("[KYUUBI #3385] generate executor pod name prefix with user or UUID") {
     val userName1 = "/kyuubi_user+-*"
     val executorPodNamePrefix1 = SparkSQLEngine.generateExecutorPodNamePrefixForK8s(userName1)