[SPARK-50648][CORE] Cleanup zombie tasks in non-running stages when the job is cancelled

core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala

@@ -2187,7 +2187,8 @@ private[spark] class DAGScheduler(
                   log"we will roll back and rerun below stages which include itself and all its " +
                   log"indeterminate child stages: ${MDC(STAGES, rollingBackStages)}")
               }
-
+              failedStage.markResubmitInFetchFailed()
+              mapStage.markResubmitInFetchFailed()
               // We expect one executor failure to trigger many FetchFailures in rapid succession,
               // but all of those task failures can typically be handled by a single resubmission of
               // the failed stage.  We avoid flooding the scheduler's event queue with resubmit
@@ -2937,7 +2938,9 @@ private[spark] class DAGScheduler(
         } else {
           // This stage is only used by the job, so finish the stage if it is running.
           val stage = stageIdToStage(stageId)
-          if (runningStages.contains(stage)) {
+          val shouldKill = runningStages.contains(stage) ||
+            (waitingStages.contains(stage) && stage.resubmitInFetchFailed)
+          if (shouldKill) {
             try { // killAllTaskAttempts will fail if a SchedulerBackend does not implement killTask
               taskScheduler.killAllTaskAttempts(stageId, shouldInterruptTaskThread(job), reason)
               if (legacyAbortStageAfterKillTasks) {

core/src/main/scala/org/apache/spark/scheduler/Stage.scala

@@ -71,6 +71,7 @@ private[scheduler] abstract class Stage(
   /** The ID to use for the next new attempt for this stage. */
   private var nextAttemptId: Int = 0
   private[scheduler] def getNextAttemptId: Int = nextAttemptId
+  private[scheduler] var _resubmitInFetchFailed: Boolean = false
 
   val name: String = callSite.shortForm
   val details: String = callSite.longForm
@@ -96,6 +97,12 @@ private[scheduler] abstract class Stage(
     failedAttemptIds.clear()
   }
 
+  private[scheduler] def resubmitInFetchFailed: Boolean = _resubmitInFetchFailed
+
+  private[scheduler] def markResubmitInFetchFailed() : Unit = {
+    _resubmitInFetchFailed = true
+  }
+
   /** Creates a new attempt for this stage by creating a new StageInfo with a new attempt ID. */
   def makeNewStageAttempt(
       numPartitionsToCompute: Int,

core/src/main/scala/org/apache/spark/scheduler/TaskScheduler.scala

@@ -121,4 +121,6 @@ private[spark] trait TaskScheduler {
    */
   def applicationAttemptId(): Option[String]
 
+
+  def hasRunningTasks(stageId: Int): Boolean
 }

core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala

@@ -1212,6 +1212,16 @@ private[spark] class TaskSchedulerImpl(
 
   override def applicationAttemptId(): Option[String] = backend.applicationAttemptId()
 
+  override def hasRunningTasks(stageId: Int): Boolean = synchronized {
+    var hasRunningTasks = false
+    taskSetsByStageIdAndAttempt.get(stageId).foreach { attempts =>
+      attempts.foreach { case (_, tsm) =>
+        hasRunningTasks = hasRunningTasks || tsm.runningTasksSet.nonEmpty
+      }
+    }
+    hasRunningTasks
+  }
+
   // exposed for testing
   private[scheduler] def taskSetManagerForAttempt(
       stageId: Int,

core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala

@@ -225,6 +225,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
     override def executorLost(executorId: String, reason: ExecutorLossReason): Unit = {}
     override def workerRemoved(workerId: String, host: String, message: String): Unit = {}
     override def applicationAttemptId(): Option[String] = None
+    override def hasRunningTasks(stageId: Int): Boolean = false
     override def executorDecommission(
       executorId: String,
       decommissionInfo: ExecutorDecommissionInfo): Unit = {
@@ -941,6 +942,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
       override def executorLost(executorId: String, reason: ExecutorLossReason): Unit = {}
       override def workerRemoved(workerId: String, host: String, message: String): Unit = {}
       override def applicationAttemptId(): Option[String] = None
+      override def hasRunningTasks(stageId: Int): Boolean = false
       override def executorDecommission(
         executorId: String,
         decommissionInfo: ExecutorDecommissionInfo): Unit = {}
@@ -2248,7 +2250,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
     // original result task 1.0 succeed
     runEvent(makeCompletionEvent(taskSets(1).tasks(1), Success, 42))
     sc.listenerBus.waitUntilEmpty()
-    assert(completedStage === List(0, 1, 1, 0))
+    assert(completedStage === List(0, 1, 1, 0, 1))
     assert(scheduler.activeJobs.isEmpty)
   }
 

core/src/test/scala/org/apache/spark/scheduler/ExternalClusterManagerSuite.scala

@@ -95,6 +95,7 @@ private class DummyTaskScheduler extends TaskScheduler {
   override def executorLost(executorId: String, reason: ExecutorLossReason): Unit = {}
   override def workerRemoved(workerId: String, host: String, message: String): Unit = {}
   override def applicationAttemptId(): Option[String] = None
+  override def hasRunningTasks(stageId: Int): Boolean = false
   def executorHeartbeatReceived(
       execId: String,
       accumUpdates: Array[(Long, Seq[AccumulatorV2[_, _]])],

sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala

@@ -938,8 +938,12 @@ class AdaptiveQueryExecSuite
           val error = intercept[SparkException] {
             joined.collect()
           }
-          assert((Seq(error) ++ Option(error.getCause) ++ error.getSuppressed()).exists(
-            e => e.getMessage() != null && e.getMessage().contains("coalesce test error")))
+          val errorMessages = (Seq(error) ++ Option(error.getCause) ++ error.getSuppressed())
+            .filter(e => e.getMessage != null).map(e => e.getMessage)
+          assert(errorMessages.exists(
+            e => e.contains("coalesce test error")),
+            s"Error messages should contain `coalesce test error`, " +
+              s"error messages: $errorMessages")
 
           val adaptivePlan = joined.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]