Fix race issue

Author: sarutak

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala

@@ -193,6 +193,11 @@ private[connect] class ExecuteThreadRunner(executeHolder: ExecuteHolder) extends
       return
     }
 
+    // SPARK-53339: Post the Started event here, right after the CAS succeeds, to ensure that
+    // postStarted() is never called when interrupt() has already transitioned the state to
+    // interrupted. This eliminates the race between postStarted() and interrupt().
+    executeHolder.eventsManager.postStarted()
+
     // `withSession` ensures that session-specific artifacts (such as JARs and class files) are
     // available during processing.
     executeHolder.sessionHolder.withSession { session =>

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteEventsManager.scala

@@ -248,8 +248,11 @@ case class ExecuteEventsManager(executeHolder: ExecuteHolder, clock: Clock) {
    * Post @link org.apache.spark.sql.connect.service.SparkListenerConnectOperationCanceled.
    */
   def postCanceled(): Unit = {
+    // SPARK-53339: Pending is included to handle the case where interrupt() is called before
+    // postStarted() transitions the status from Pending to Started.
     assertStatus(
       List(
+        ExecuteStatus.Pending,
         ExecuteStatus.Started,
         ExecuteStatus.Analyzed,
         ExecuteStatus.ReadyForExecution,
@@ -269,8 +272,11 @@ case class ExecuteEventsManager(executeHolder: ExecuteHolder, clock: Clock) {
    *   The message of the error thrown during the request.
    */
   def postFailed(errorMessage: String): Unit = {
+    // SPARK-53339: Pending is included to handle the case where postStarted() itself throws
+    // an exception (e.g., session state check failure) before transitioning from Pending.
     assertStatus(
       List(
+        ExecuteStatus.Pending,
         ExecuteStatus.Started,
         ExecuteStatus.Analyzed,
         ExecuteStatus.ReadyForExecution,

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala

@@ -191,7 +191,6 @@ private[connect] class SparkConnectExecutionManager() extends Logging {
       responseObserver: StreamObserver[proto.ExecutePlanResponse]): ExecuteHolder = {
     val executeHolder = createExecuteHolder(executeKey, request, sessionHolder)
     try {
-      executeHolder.eventsManager.postStarted()
       executeHolder.start()
     } catch {
       // Errors raised before the execution holder has finished spawning a thread are considered

sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/ExecuteEventsManagerSuite.scala

@@ -138,6 +138,36 @@ class ExecuteEventsManagerSuite
         .isInstanceOf[SparkListenerConnectOperationCanceled])
   }
 
+  test("SPARK-53339: post canceled from Pending state") {
+    val events = setupEvents(ExecuteStatus.Pending)
+    events.postCanceled()
+    assert(events.status == ExecuteStatus.Canceled)
+    assert(events.terminationReason.contains(TerminationReason.Canceled))
+  }
+
+  test("SPARK-53339: post failed from Pending state") {
+    val events = setupEvents(ExecuteStatus.Pending)
+    events.postFailed(DEFAULT_ERROR)
+    assert(events.status == ExecuteStatus.Failed)
+    assert(events.terminationReason.contains(TerminationReason.Failed))
+  }
+
+  test("SPARK-53339: Pending to Canceled to Closed transition") {
+    val events = setupEvents(ExecuteStatus.Pending)
+    events.postCanceled()
+    events.postClosed()
+    assert(events.status == ExecuteStatus.Closed)
+    assert(events.terminationReason.contains(TerminationReason.Canceled))
+  }
+
+  test("SPARK-53339: Pending to Failed to Closed transition") {
+    val events = setupEvents(ExecuteStatus.Pending)
+    events.postFailed(DEFAULT_ERROR)
+    events.postClosed()
+    assert(events.status == ExecuteStatus.Closed)
+    assert(events.terminationReason.contains(TerminationReason.Failed))
+  }
+
   test("SPARK-43923: post failed") {
     val events = setupEvents(ExecuteStatus.Started)
     events.postFailed(DEFAULT_ERROR)