(spark) branch master updated: [SPARK-58887][CORE] Fix barrier job to be cancellable during slot check retries

[do...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dongjoon-hyun pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new 4998b7769ba4 [SPARK-58887][CORE] Fix barrier job to be cancellable during slot check retries
4998b7769ba4 is described below

commit 4998b7769ba49642b9167903b503336d6c8f02d2
Author: Dongjoon Hyun <do...@apache.org>
AuthorDate: Wed Aug 19 21:54:18 2026 -0700

    [SPARK-58887][CORE] Fix barrier job to be cancellable during slot check retries
    
    ### What changes were proposed in this pull request?
    
    This PR makes a barrier job cancellable while its max-concurrent-tasks slot check is being
    retried, by tracking the deferred submission in `DAGScheduler`:
    
    - Add `deferredBarrierJobs`, a map from job id to the deferred job's `JobListener` and its
      submission-time `Properties`.
    - In the `BarrierJobSlotsNumberCheckFailed` retry branch of `handleJobSubmitted`, register the
      deferred job in the map before scheduling the re-post.
    - On cancellation (`handleJobCancellation`), fail the deferred listener directly with
      `sparkJobCancelled` and swap the map entry to a `deferredJobCancelledMarker` tombstone (also
      dropping the job's entry in `barrierJobIdToNumTasksCheckFailures` so the failure counter does
      not leak). The pending re-post always fires; `handleJobSubmitted` drops the re-posted
      submission when it finds the marker, so the cancel-vs-re-post decision is made entirely on
      the event loop and races nothing.
    - Also clean up the cancelled job's partial stage registrations: stage creation may have
      registered ancestor stages (e.g. an ordinary shuffle upstream of the barrier stage) before
      the slot check threw. The stage-unregistration part of
      `cleanupStateForJobAndIndependentStages` is extracted into a job-id-keyed
      `cleanupStagesForJob` and invoked on the deferred cancellation, so the abandoned
      registrations cannot break a later cancellation of the same job id (which would otherwise
      reach `jobIdToActiveJob(jobId)` and crash the event loop) or pin stages other jobs do not
      need.
    - Make `handleJobGroupCancelled` and `handleJobTagCancelled` also match deferred barrier jobs
      by the properties captured at submission, so `cancelJobGroup` / `cancelJobsWithTag` (the
      paths used by Spark Connect interrupt, Thrift server cancel, and streaming stop) reach them
      too.
    - Report deferred cancellations through the tag-cancellation result: the promise payload of
      `SparkContext.cancelJobsWithTagWithFuture` changes from `Seq[ActiveJob]` to a new
      `private[spark] case class CancelledJobInfo(jobId, properties)` covering both active and
      deferred jobs, so classic `SparkSession.interruptTag` / `interruptOperation` /
      `interruptAll` return the SQL execution ids of deferred barrier jobs they cancelled. The
      erased signature is unchanged and all touched APIs are `private[spark]`.
    - Drain the map in `doCancelAllJobs` (failing each deferred job with the cancel-all reason) and
      in `cleanUpAfterSchedulerStop` (failing each with "SparkContext was shut down").
    
    ### Why are the changes needed?
    
    When a barrier job fails the max-concurrent-tasks slot check, `handleJobSubmitted` re-posts the
    `JobSubmitted` event every `spark.scheduler.barrier.maxConcurrentTasksCheck.interval` (15s) up
    to `spark.scheduler.barrier.maxConcurrentTasksCheck.maxFailures` (40) times, an ~10 minute
    window by default. During that window the job has no `ActiveJob` (`createResultStage` aborts
    with the exception before one is created), so:
    
    - `sc.cancelJob(jobId)`, `sc.cancelAllJobs()`, `sc.cancelJobGroup(...)` and
      `sc.cancelJobsWithTag(...)` are silent no-ops, and
    - no `SparkListenerJobStart` was posted, so the job is invisible in the UI.
    
    The user has no way to cancel the job until the retries are exhausted.
    
    ### Does this PR introduce _any_ user-facing change?
    
    Yes, it is a bug fix. Previously, cancelling a barrier job during its slot-check retry window
    was silently ignored and the job kept retrying; now the job fails immediately with the usual
    job-cancelled error, and the pending re-post is dropped.
    
    ### How was this patch tested?
    
    Added six tests to `DAGSchedulerSuite`, all submitting a barrier job whose slot check fails on
    the `local[2]` backend so the job is deferred:
    
    - `SPARK-58887: a barrier job can be cancelled while its slot check is being retried` —
      cancels the job during the retry window and verifies the listener fails immediately with the
      cancellation error, then simulates the pending re-post arriving after the cancellation and
      verifies it is dropped instead of resurrecting the job, with all internal maps cleaned up.
    - `SPARK-58887: cancelAllJobs also fails a barrier job deferred for a slot-check retry`
    - `SPARK-58887: cancelJobGroup also fails a barrier job deferred for a slot-check retry`
    - `SPARK-58887: cancelJobsWithTag also fails a deferred barrier job`
    - `SPARK-58887: tag cancellation reports a deferred barrier job in its promise` — verifies the
      deferred job's id and submission-time properties are reported through the
      `cancelJobsWithTagWithFuture` promise.
    - `SPARK-58887: cancelling a deferred barrier job drops its partial stage registrations` —
      uses an ordinary-shuffle -> barrier-shuffle -> result graph so an ancestor stage is
      registered before the slot check throws, and verifies the cancellation drops those
      registrations and a later cancellation of the same job id is a harmless no-op.
    
    Also ran the full `DAGSchedulerSuite`, `BarrierStageOnSubmittedSuite`, and
    `JobCancellationSuite` (256 tests, all passed), plus
    `SparkSessionJobTaggingAndCancellationSuite` in `sql/core` (6 tests, all passed) since the
    tag-cancellation future type changed.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Claude Fable 5
    
    Closes #58138 from dongjoon-hyun/SPARK-58887.
    
    Authored-by: Dongjoon Hyun <do...@apache.org>
    Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
 .../main/scala/org/apache/spark/SparkContext.scala |   9 +-
 .../org/apache/spark/scheduler/DAGScheduler.scala  | 156 ++++++++++++++++++---
 .../apache/spark/scheduler/DAGSchedulerEvent.scala |   2 +-
 .../apache/spark/scheduler/DAGSchedulerSuite.scala | 144 +++++++++++++++++++
 4 files changed, 284 insertions(+), 27 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index 1d27dea2b3f4..72830a565f80 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -2739,16 +2739,17 @@ class SparkContext(config: SparkConf) extends Logging {
    *
    * @param tag The tag to be cancelled. Cannot contain ',' (comma) character.
    * @param reason reason for cancellation.
-   * @return A future with [[ActiveJob]]s, allowing extraction of information such as Job ID and
-   *   tags.
+   * @return A future with the cancelled jobs' [[CancelledJobInfo]], allowing extraction of
+   *   information such as Job ID and tags. Covers active jobs and barrier jobs cancelled while
+   *   deferred for their slot-check retry (which have no [[ActiveJob]]).
    */
   private[spark] def cancelJobsWithTagWithFuture(
       tag: String,
-      reason: String): Future[Seq[ActiveJob]] = {
+      reason: String): Future[Seq[CancelledJobInfo]] = {
     SparkContext.throwIfInvalidTag(tag)
     assertNotStopped()
 
-    val cancelledJobs = Promise[Seq[ActiveJob]]()
+    val cancelledJobs = Promise[Seq[CancelledJobInfo]]()
     dagScheduler.cancelJobsWithTag(tag, Some(reason), Some(cancelledJobs))
     cancelledJobs.future
   }
diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
index 652d38b7a4e8..cdc65c035a60 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
@@ -406,6 +406,35 @@ private[spark] class DAGScheduler(
    */
   private[scheduler] val barrierJobIdToNumTasksCheckFailures = new ConcurrentHashMap[Int, Int]
 
+  /**
+   * The barrier jobs deferred while their max concurrent tasks check is being retried (the
+   * submission is re-posted on a timer), keyed by job id. A deferred job is registered nowhere
+   * else, so this is what lets a cancellation fail it: the cancellation swaps in
+   * `deferredJobCancelledMarker` and `handleJobSubmitted` drops the re-posted submission when
+   * it finds the marker.
+   */
+  private[scheduler] val deferredBarrierJobs =
+    new ConcurrentHashMap[Int, DAGScheduler.DeferredBarrierJob]
+
+  /**
+   * Marker left in `deferredBarrierJobs` when a deferred job is cancelled. The entry is
+   * replaced rather than removed so that the decision to drop the pending re-post is made on
+   * the event loop (in `handleJobSubmitted`), atomically with respect to the cancellation; the
+   * re-post itself always fires and is what finally clears the entry.
+   */
+  private val deferredJobCancelledMarker = DAGScheduler.DeferredBarrierJob(null, null)
+
+  /** (id, properties) of the deferred barrier jobs whose submission-time properties match `p`. */
+  private def deferredJobsMatching(p: Properties => Boolean): Seq[(Int, Properties)] = {
+    val matched = mutable.ArrayBuffer[(Int, Properties)]()
+    deferredBarrierJobs.forEach { (jobId, deferred) =>
+      if ((deferred ne deferredJobCancelledMarker) && p(deferred.properties)) {
+        matched += ((jobId, deferred.properties))
+      }
+    }
+    matched.toSeq
+  }
+
   /**
    * Time in seconds to wait between a max concurrent tasks check failure and the next check.
    */
@@ -1329,24 +1358,25 @@ private[spark] class DAGScheduler(
   }
 
   /**
-   * Removes state for job and any stages that are not needed by any other job.  Does not
-   * handle cancelling tasks or notifying the SparkListener about finished jobs/stages/tasks.
-   *
-   * @param job The job whose state to cleanup.
+   * Removes the given job from every stage registered for it, unregistering each stage that no
+   * other job needs. Unlike [[cleanupStateForJobAndIndependentStages]] this does not require an
+   * `ActiveJob`: stage creation may register ancestor stages before a barrier slot check throws
+   * (see SPARK-58887), and a job cancelled while deferred for the retry must still drop those
+   * registrations.
    */
-  private def cleanupStateForJobAndIndependentStages(job: ActiveJob): Unit = {
-    val registeredStages = jobIdToStageIds.get(job.jobId)
+  private def cleanupStagesForJob(jobId: Int): Unit = {
+    val registeredStages = jobIdToStageIds.get(jobId)
     if (registeredStages.isEmpty || registeredStages.get.isEmpty) {
-      logError(log"No stages registered for job ${MDC(JOB_ID, job.jobId)}")
+      logError(log"No stages registered for job ${MDC(JOB_ID, jobId)}")
     } else {
       stageIdToStage.filter {
         case (stageId, _) => registeredStages.get.contains(stageId)
       }.foreach {
         case (stageId, stage) =>
           val jobSet = stage.jobIds
-          if (!jobSet.contains(job.jobId)) {
+          if (!jobSet.contains(jobId)) {
             // scalastyle:off line.size.limit
-            logError(log"Job ${MDC(JOB_ID, job.jobId)} not registered for stage ${MDC(STAGE_ID, stageId)} even though that stage was registered for the job")
+            logError(log"Job ${MDC(JOB_ID, jobId)} not registered for stage ${MDC(STAGE_ID, stageId)} even though that stage was registered for the job")
             // scalastyle:on
           } else {
             def removeStage(stageId: Int): Unit = {
@@ -1391,14 +1421,24 @@ private[spark] class DAGScheduler(
                 .format(stageId, stageIdToStage.size))
             }
 
-            jobSet -= job.jobId
+            jobSet -= jobId
             if (jobSet.isEmpty) { // no other job needs this stage
               removeStage(stageId)
             }
           }
       }
     }
-    jobIdToStageIds -= job.jobId
+    jobIdToStageIds -= jobId
+  }
+
+  /**
+   * Removes state for job and any stages that are not needed by any other job.  Does not
+   * handle cancelling tasks or notifying the SparkListener about finished jobs/stages/tasks.
+   *
+   * @param job The job whose state to cleanup.
+   */
+  private def cleanupStateForJobAndIndependentStages(job: ActiveJob): Unit = {
+    cleanupStagesForJob(job.jobId)
     jobIdToActiveJob -= job.jobId
     activeJobs -= job
     job.finalStage match {
@@ -1620,7 +1660,7 @@ private[spark] class DAGScheduler(
   def cancelJobsWithTag(
       tag: String,
       reason: Option[String],
-      cancelledJobs: Option[Promise[Seq[ActiveJob]]]): Unit = {
+      cancelledJobs: Option[Promise[Seq[CancelledJobInfo]]]): Unit = {
     SparkContext.throwIfInvalidTag(tag)
     logInfo(log"Asked to cancel jobs with tag ${MDC(TAG, tag)}")
     eventProcessLoop.post(JobTagCancelled(tag, reason, cancelledJobs))
@@ -1652,6 +1692,11 @@ private[spark] class DAGScheduler(
     // the caller supplied one.
     val updatedReason = reason.getOrElse(DAGScheduler.DEFAULT_CANCEL_ALL_JOBS_REASON)
     runningStages.map(_.firstJobId).foreach(handleJobCancellation(_, Option(updatedReason)))
+    // Also fail the barrier jobs deferred for a slot-check retry: they are registered nowhere
+    // else, and their pending re-posts are dropped once the entries are marked cancelled.
+    deferredBarrierJobs.keySet().forEach { jobId =>
+      handleJobCancellation(jobId, Option(updatedReason))
+    }
     activeJobs.clear() // These should already be empty by this point,
     jobIdToActiveJob.clear() // but just in case we lost track of some jobs...
   }
@@ -1926,11 +1971,16 @@ private[spark] class DAGScheduler(
         _.getProperty(SparkContext.SPARK_JOB_GROUP_ID) == groupId
       }
     }
-    if (activeInGroup.isEmpty && !cancelFutureJobs) {
+    // A barrier job deferred for a slot-check retry is not in `activeJobs` yet, so match it by
+    // the properties captured at submission.
+    val deferredInGroup = deferredJobsMatching { properties =>
+      Option(properties).exists(_.getProperty(SparkContext.SPARK_JOB_GROUP_ID) == groupId)
+    }
+    if (activeInGroup.isEmpty && deferredInGroup.isEmpty && !cancelFutureJobs) {
       logWarning(log"Failed to cancel job group ${MDC(GROUP_ID, groupId)}. " +
         log"Cannot find active jobs for it.")
     }
-    val jobIds = activeInGroup.map(_.jobId)
+    val jobIds = activeInGroup.map(_.jobId) ++ deferredInGroup.map(_._1)
     val updatedReason = reason.getOrElse("part of cancelled job group %s".format(groupId))
     jobIds.foreach(handleJobCancellation(_, Option(updatedReason)))
   }
@@ -1938,19 +1988,32 @@ private[spark] class DAGScheduler(
   private[scheduler] def handleJobTagCancelled(
       tag: String,
       reason: Option[String],
-      cancelledJobs: Option[Promise[Seq[ActiveJob]]]): Unit = {
+      cancelledJobs: Option[Promise[Seq[CancelledJobInfo]]]): Unit = {
     // Cancel all jobs that have all provided tags.
     // First finds all active jobs with this group id, and then kill stages for them.
     val jobsToBeCancelled = activeJobs.filter { activeJob =>
-      Option(activeJob.properties).exists { properties =>
-        Option(properties.getProperty(SparkContext.SPARK_JOB_TAGS)).getOrElse("")
-          .split(SparkContext.SPARK_JOB_TAGS_SEP).filter(!_.isEmpty).toSet.contains(tag)
-      }
+      hasJobTag(activeJob.properties, tag)
     }
+    // A barrier job deferred for a slot-check retry is not in `activeJobs` yet, so match it by
+    // the properties captured at submission.
+    val deferredTagged = deferredJobsMatching(hasJobTag(_, tag))
     val updatedReason =
       reason.getOrElse("part of cancelled job tags %s".format(tag))
-    jobsToBeCancelled.map(_.jobId).foreach(handleJobCancellation(_, Option(updatedReason)))
-    cancelledJobs.map(_.success(jobsToBeCancelled.toSeq))
+    (jobsToBeCancelled.map(_.jobId) ++ deferredTagged.map(_._1))
+      .foreach(handleJobCancellation(_, Option(updatedReason)))
+    // Report the deferred jobs too: they have no ActiveJob, but consumers (e.g. classic
+    // SparkSession.interruptTag) read the SQL execution id from the properties.
+    cancelledJobs.map(_.success(
+      jobsToBeCancelled.toSeq.map(job => CancelledJobInfo(job.jobId, job.properties)) ++
+        deferredTagged.map { case (jobId, properties) => CancelledJobInfo(jobId, properties) }))
+  }
+
+  /** Whether the job properties carry the given job tag. */
+  private def hasJobTag(properties: Properties, tag: String): Boolean = {
+    Option(properties).exists { props =>
+      Option(props.getProperty(SparkContext.SPARK_JOB_TAGS)).getOrElse("")
+        .split(SparkContext.SPARK_JOB_TAGS_SEP).filter(!_.isEmpty).toSet.contains(tag)
+    }
   }
 
   private[scheduler] def handleBeginEvent(task: Task[_], taskInfo: TaskInfo): Unit = {
@@ -2005,6 +2068,15 @@ private[spark] class DAGScheduler(
       }
       listenerBus.post(SparkListenerJobEnd(job.jobId, clock.getTimeMillis(), JobFailed(error)))
     }
+    // Also complete the waiters of the barrier jobs deferred for a slot-check retry: they are
+    // in `activeJobs` above only once re-processed, which will never happen now.
+    deferredBarrierJobs.forEach { (jobId, deferred) =>
+      if (deferred ne deferredJobCancelledMarker) {
+        deferred.listener.jobFailed(
+          new SparkException(s"Job $jobId cancelled because SparkContext was shut down"))
+      }
+    }
+    deferredBarrierJobs.clear()
   }
 
   private[scheduler] def handleGetTaskResult(taskInfo: TaskInfo): Unit = {
@@ -2032,6 +2104,11 @@ private[spark] class DAGScheduler(
       listener: JobListener,
       artifacts: JobArtifactSet,
       properties: Properties): Unit = {
+    // The job is being (re-)processed, so it is no longer merely deferred. The marker means it
+    // was cancelled while deferred: its listener has already been failed, so drop this re-post.
+    if (deferredBarrierJobs.remove(jobId) eq deferredJobCancelledMarker) {
+      return
+    }
     // If this job belongs to a cancelled job group, skip running it
     val jobGroupIdOpt = Option(properties).map(_.getProperty(SparkContext.SPARK_JOB_GROUP_ID))
     if (jobGroupIdOpt.exists(cancelledJobGroups.contains(_))) {
@@ -2103,6 +2180,7 @@ private[spark] class DAGScheduler(
           log"more times")
 
         if (numCheckFailures <= maxFailureNumTasksCheck) {
+          deferredBarrierJobs.put(jobId, DAGScheduler.DeferredBarrierJob(listener, properties))
           messageScheduler.schedule(
             new Runnable {
               override def run(): Unit = eventProcessLoop.post(JobSubmitted(jobId, finalRDD, func,
@@ -4284,7 +4362,25 @@ private[spark] class DAGScheduler(
   }
 
   private[scheduler] def handleJobCancellation(jobId: Int, reason: Option[String]): Unit = {
-    if (!jobIdToStageIds.contains(jobId)) {
+    val deferred = deferredBarrierJobs.get(jobId)
+    if (deferred != null) {
+      if (deferred ne deferredJobCancelledMarker) {
+        // A barrier job deferred for a slot-check retry is registered nowhere else, so fail its
+        // listener directly. Leave the marker in place instead of removing the entry: the
+        // pending re-post always fires, and handleJobSubmitted drops it on finding the marker.
+        deferredBarrierJobs.put(jobId, deferredJobCancelledMarker)
+        barrierJobIdToNumTasksCheckFailures.remove(jobId)
+        // Stage creation may have registered ancestor stages (e.g. an ordinary shuffle upstream
+        // of the barrier stage) before the slot check threw. Drop this job from them, and
+        // unregister the stages no other job needs, so the abandoned registrations cannot break
+        // a later cancellation of this job id or pin the stages.
+        if (jobIdToStageIds.contains(jobId)) {
+          cleanupStagesForJob(jobId)
+        }
+        deferred.listener.jobFailed(
+          SparkCoreErrors.sparkJobCancelled(jobId, reason.getOrElse(""), null))
+      }
+    } else if (!jobIdToStageIds.contains(jobId)) {
       logDebug("Trying to cancel unregistered job " + jobId)
     } else {
       failJobAndIndependentStages(
@@ -4765,8 +4861,24 @@ private[spark] object DAGScheduler {
   // Fallback reason used when a context-wide cancellation does not supply a more specific one.
   // Kept as the historical wording so existing log and error consumers are unaffected.
   val DEFAULT_CANCEL_ALL_JOBS_REASON = "as part of cancellation of all jobs"
+
+  /**
+   * A barrier job deferred while its max concurrent tasks check is being retried, tracked in
+   * `deferredBarrierJobs`. The submission-time properties are kept so group/tag cancellations
+   * can match the job.
+   */
+  private[scheduler] case class DeferredBarrierJob(listener: JobListener, properties: Properties)
 }
 
+/**
+ * Metadata of a job cancelled by a tag cancellation, reported through the promise of
+ * `SparkContext.cancelJobsWithTagWithFuture`. Not restricted to jobs with an `ActiveJob`: a
+ * barrier job cancelled while deferred for its slot-check retry is reported too, and consumers
+ * (e.g. classic `SparkSession.interruptTag`) read the SQL execution id from the submission-time
+ * properties.
+ */
+private[spark] case class CancelledJobInfo(jobId: Int, properties: Properties)
+
 /**
  * Thrown when a job uses a pipelined-shuffle idiom that is not supported (fan-out, a barrier /
  * indeterminate / checksum-retry / push-merge producer, a reliable checkpoint in a member's chain,
diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala
index 44811a721483..20905ba58ba1 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala
@@ -74,7 +74,7 @@ private[scheduler] case class JobGroupCancelled(
 private[scheduler] case class JobTagCancelled(
     tagName: String,
     reason: Option[String],
-    cancelledJobs: Option[Promise[Seq[ActiveJob]]]) extends DAGSchedulerEvent
+    cancelledJobs: Option[Promise[Seq[CancelledJobInfo]]]) extends DAGSchedulerEvent
 
 private[scheduler] case class AllJobsCancelled(reason: Option[String] = None)
   extends DAGSchedulerEvent
diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
index d18b260d7b24..f4dcdaf96d7c 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
@@ -23,6 +23,7 @@ import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicLong, At
 
 import scala.annotation.meta.param
 import scala.collection.mutable.{ArrayBuffer, HashMap, HashSet, Map}
+import scala.concurrent.Promise
 import scala.jdk.CollectionConverters._
 import scala.language.reflectiveCalls
 import scala.util.control.NonFatal
@@ -1559,6 +1560,149 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
     assertDataStructuresEmpty()
   }
 
+  test("SPARK-58887: a barrier job can be cancelled while its slot check is being retried") {
+    // 3 barrier tasks on the local[2] backend fail the max concurrent tasks check, so the
+    // submission enters the retry window during which the job is registered nowhere but
+    // `deferredBarrierJobs`. A cancellation in that window must fail the job immediately
+    // instead of silently no-oping until the retries run out.
+    val barrierRdd = new MyRDD(sc, 3, Nil).barrier().mapPartitions(iter => iter)
+    val failure = new AtomicReference[Exception]()
+    val failListener = new JobListener {
+      override def taskSucceeded(index: Int, result: Any): Unit = {}
+      override def jobFailed(exception: Exception): Unit = failure.set(exception)
+    }
+    val jobId = submit(barrierRdd, Array(0, 1, 2), listener = failListener)
+    assert(failure.get() === null, "a barrier job in its slot-check retry window must wait")
+    assert(scheduler.deferredBarrierJobs.containsKey(jobId))
+    assert(scheduler.barrierJobIdToNumTasksCheckFailures.containsKey(jobId))
+    cancel(jobId)
+    assert(failure.get() !== null, "cancelling a deferred barrier job must fail its listener")
+    assert(failure.get().getMessage.contains(s"Job $jobId cancelled"),
+      "the listener must fail with the cancellation error, not the exhausted-retries one")
+    assert(!scheduler.barrierJobIdToNumTasksCheckFailures.containsKey(jobId),
+      "a cancelled deferral must not leak its slot-check failure count")
+    // The pending re-post fires regardless of the cancellation; simulate its arrival and check
+    // it is dropped instead of resurrecting the cancelled job.
+    runEvent(JobSubmitted(jobId, barrierRdd, jobComputeFunc, Array(0, 1, 2), CallSite("", ""),
+      failListener, JobArtifactSet.getActiveOrDefault(sc), null))
+    assert(scheduler.deferredBarrierJobs.isEmpty,
+      "a re-post arriving after the cancellation must be dropped, not re-deferred")
+    assert(scheduler.barrierJobIdToNumTasksCheckFailures.isEmpty)
+    assertDataStructuresEmpty()
+  }
+
+  test("SPARK-58887: cancelAllJobs also fails a barrier job deferred for a slot-check retry") {
+    val barrierRdd = new MyRDD(sc, 3, Nil).barrier().mapPartitions(iter => iter)
+    val failure = new AtomicReference[Exception]()
+    val failListener = new JobListener {
+      override def taskSucceeded(index: Int, result: Any): Unit = {}
+      override def jobFailed(exception: Exception): Unit = failure.set(exception)
+    }
+    val jobId = submit(barrierRdd, Array(0, 1, 2), listener = failListener)
+    assert(failure.get() === null)
+    assert(scheduler.deferredBarrierJobs.containsKey(jobId))
+    runEvent(AllJobsCancelled())
+    assert(failure.get() !== null, "cancelAllJobs must fail a deferred barrier job's listener")
+    assert(failure.get().getMessage.contains(s"Job $jobId cancelled"))
+    assert(!scheduler.barrierJobIdToNumTasksCheckFailures.containsKey(jobId))
+    assertDataStructuresEmpty()
+  }
+
+  test("SPARK-58887: cancelJobGroup also fails a barrier job deferred for a slot-check retry") {
+    val barrierRdd = new MyRDD(sc, 3, Nil).barrier().mapPartitions(iter => iter)
+    val failure = new AtomicReference[Exception]()
+    val failListener = new JobListener {
+      override def taskSucceeded(index: Int, result: Any): Unit = {}
+      override def jobFailed(exception: Exception): Unit = failure.set(exception)
+    }
+    val props = new Properties()
+    props.setProperty(SparkContext.SPARK_JOB_GROUP_ID, "deferredGroup")
+    val jobId = submit(barrierRdd, Array(0, 1, 2), listener = failListener, properties = props)
+    assert(failure.get() === null)
+    assert(scheduler.deferredBarrierJobs.containsKey(jobId))
+    runEvent(JobGroupCancelled("deferredGroup", cancelFutureJobs = false, None))
+    assert(failure.get() !== null, "cancelJobGroup must fail a deferred barrier job's listener")
+    assert(failure.get().getMessage.contains(s"Job $jobId cancelled"))
+    assert(!scheduler.barrierJobIdToNumTasksCheckFailures.containsKey(jobId))
+    assertDataStructuresEmpty()
+  }
+
+  test("SPARK-58887: cancelJobsWithTag also fails a deferred barrier job") {
+    val barrierRdd = new MyRDD(sc, 3, Nil).barrier().mapPartitions(iter => iter)
+    val failure = new AtomicReference[Exception]()
+    val failListener = new JobListener {
+      override def taskSucceeded(index: Int, result: Any): Unit = {}
+      override def jobFailed(exception: Exception): Unit = failure.set(exception)
+    }
+    val props = new Properties()
+    props.setProperty(SparkContext.SPARK_JOB_TAGS, "deferredTag")
+    val jobId = submit(barrierRdd, Array(0, 1, 2), listener = failListener, properties = props)
+    assert(failure.get() === null)
+    assert(scheduler.deferredBarrierJobs.containsKey(jobId))
+    runEvent(JobTagCancelled("deferredTag", None, None))
+    assert(failure.get() !== null,
+      "cancelJobsWithTag must fail a deferred barrier job's listener")
+    assert(failure.get().getMessage.contains(s"Job $jobId cancelled"))
+    assert(!scheduler.barrierJobIdToNumTasksCheckFailures.containsKey(jobId))
+    assertDataStructuresEmpty()
+  }
+
+  test("SPARK-58887: tag cancellation reports a deferred barrier job in its promise") {
+    val barrierRdd = new MyRDD(sc, 3, Nil).barrier().mapPartitions(iter => iter)
+    val failure = new AtomicReference[Exception]()
+    val failListener = new JobListener {
+      override def taskSucceeded(index: Int, result: Any): Unit = {}
+      override def jobFailed(exception: Exception): Unit = failure.set(exception)
+    }
+    val props = new Properties()
+    props.setProperty(SparkContext.SPARK_JOB_TAGS, "reportedTag")
+    val jobId = submit(barrierRdd, Array(0, 1, 2), listener = failListener, properties = props)
+    assert(scheduler.deferredBarrierJobs.containsKey(jobId))
+    val cancelledJobs = Promise[Seq[CancelledJobInfo]]()
+    runEvent(JobTagCancelled("reportedTag", None, Some(cancelledJobs)))
+    assert(failure.get() !== null)
+    // The handler completes the promise synchronously; the deferred job must be reported with
+    // its submission-time properties (e.g. for SQL execution id extraction), even though it
+    // never had an ActiveJob.
+    val reported = cancelledJobs.future.value.get.get
+    assert(reported.map(_.jobId) === Seq(jobId))
+    assert(reported.head.properties.getProperty(SparkContext.SPARK_JOB_TAGS) === "reportedTag")
+    assertDataStructuresEmpty()
+  }
+
+  test("SPARK-58887: cancelling a deferred barrier job drops its partial stage registrations") {
+    // An ordinary shuffle upstream of the barrier stage is created and registered before the
+    // barrier slot check throws, so the deferred job is NOT registered nowhere. Cancelling it
+    // must drop those registrations too, or a later cancellation of the same job id finds
+    // jobIdToStageIds populated without an ActiveJob and crashes the event loop.
+    val ordinaryRdd = new MyRDD(sc, 2, Nil)
+    val ordinaryDep = new ShuffleDependency(ordinaryRdd, new HashPartitioner(3))
+    val barrierRdd = new MyRDD(sc, 3, List(ordinaryDep), tracker = mapOutputTracker)
+      .barrier().mapPartitions(iter => iter)
+    val barrierDep = new ShuffleDependency(barrierRdd, new HashPartitioner(2))
+    val resultRdd = new MyRDD(sc, 2, List(barrierDep), tracker = mapOutputTracker)
+    val failure = new AtomicReference[Exception]()
+    val failListener = new JobListener {
+      override def taskSucceeded(index: Int, result: Any): Unit = {}
+      override def jobFailed(exception: Exception): Unit = failure.set(exception)
+    }
+    val jobId = submit(resultRdd, Array(0, 1), listener = failListener)
+    assert(failure.get() === null)
+    assert(scheduler.deferredBarrierJobs.containsKey(jobId))
+    assert(scheduler.jobIdToStageIds.contains(jobId),
+      "the ordinary ancestor stage must have been registered before the slot check threw")
+    cancel(jobId)
+    assert(failure.get() !== null)
+    assert(!scheduler.jobIdToStageIds.contains(jobId),
+      "cancelling a deferred job must drop its partial stage registrations")
+    // The pending re-post arrives and is dropped; cancelling once more must then be a harmless
+    // no-op instead of tripping over leftover registrations without an ActiveJob.
+    runEvent(JobSubmitted(jobId, resultRdd, jobComputeFunc, Array(0, 1), CallSite("", ""),
+      failListener, JobArtifactSet.getActiveOrDefault(sc), null))
+    cancel(jobId)
+    assertDataStructuresEmpty()
+  }
+
   test("Fail the job if a barrier ResultTask failed") {
     val shuffleMapRdd = new MyRDD(sc, 2, Nil)
     val shuffleDep = new ShuffleDependency(shuffleMapRdd, new HashPartitioner(2))


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org