fix concurrent promise issue

Author: summaryzb

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

@@ -51,11 +51,7 @@ private[spark] class JobWaiter[T](
    * all the tasks belonging to this job, it will fail this job with a SparkException.
    */
   def cancel(reason: Option[String] = None, quiet: Boolean = false): Unit = {
-    if (quiet) {
-      dagScheduler.cancelJob(jobId, reason, quiet)
-    } else {
-      dagScheduler.cancelJob(jobId, None, quiet)
-    }
+    dagScheduler.cancelJob(jobId, reason, quiet)
   }
 
   /**

sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala

@@ -78,8 +78,13 @@ trait ShuffleExchangeLike extends Exchange {
   private[sql] // Exposed for testing
   val futureAction = new AtomicReference[Option[FutureAction[MapOutputStatistics]]](None)
 
+  @volatile
   @transient
-  private val isCancelled: AtomicReference[Boolean] = new AtomicReference(false)
+  private var isCancelled: Boolean = false
+
+  @volatile
+  @transient
+  private var quietly: Boolean = false
 
   @transient
   private lazy val triggerFuture: java.util.concurrent.Future[Any] = {
@@ -90,7 +95,7 @@ trait ShuffleExchangeLike extends Exchange {
         executeQuery(null)
         // Submit shuffle job if not cancelled.
         this.synchronized {
-          if (isCancelled.get) {
+          if (isCancelled && !quietly) {
             promise.tryFailure(new SparkException("Shuffle cancelled."))
           } else {
             val shuffleJob = RDDOperationScope.withScope(sparkContext, nodeName, false, true) {
@@ -125,15 +130,15 @@ trait ShuffleExchangeLike extends Exchange {
    * Cancels the shuffle job with an optional reason.
    */
   final def cancelShuffleJob(reason: Option[String], quiet: Boolean): Unit = this.synchronized {
-    if (!isCancelled.get) {
-      if (quiet) {
-        // tryFailure before isCancelled.set(true) to avoid concurrent issue with triggerFuture
-        promise.tryFailure(new SparkAQEStageCancelException())
+    this.synchronized {
+      if (!isCancelled) {
+        isCancelled = true
+        if (quiet) {
+          quietly = quiet
+          promise.tryFailure(new SparkAQEStageCancelException)
+        }
+        futureAction.get().foreach(_.cancel(reason, quiet))
       }
-      isCancelled.set(true)
-      // this may take no effect, if the job has already been submitted
-      // but futureAction is not set yet, it's a best effort to cancel it
-      futureAction.get().foreach(_.cancel(reason, quiet))
     }
   }