[SPARK-39195][SQL] Spark OutputCommitCoordinator should abort stage when committed file not consistent with task status

core/src/main/scala/org/apache/spark/mapred/SparkHadoopMapRedUtil.scala

@@ -76,6 +76,8 @@ object SparkHadoopMapRedUtil extends Logging {
 
         if (canCommit) {
           performCommit()
+          outputCommitCoordinator.commitSuccess(ctx.stageId(),
+            ctx.stageAttemptNumber(), splitId, ctx.attemptNumber())
         } else {
           val message =
             s"$mrTaskAttemptID: Not committed because the driver did not authorize commit"

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

@@ -33,6 +33,12 @@ private case class AskPermissionToCommitOutput(
     partition: Int,
     attemptNumber: Int)
 
+private case class CommitOutputSuccess(
+    stage: Int,
+    stateAttempt: Int,
+    partition: Int,
+    attemptNumber: Int)
+
 /**
  * Authority that decides whether tasks can commit output to HDFS. Uses a "first committer wins"
  * policy.
@@ -55,8 +61,13 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
   // concurrently in two different attempts of the same stage.
   private case class TaskIdentifier(stageAttempt: Int, taskAttempt: Int)
 
+  // Class used to identify a committer 's status when this committer is allowed to commit task.
+  // Status is false means this committer is allowed to commit task, status is true means this
+  // committer have sent CommitOutputSuccess message to OutputCommitCoordinator.
+  private case class CommitStatus(taskIdent: TaskIdentifier, status: Boolean)
+
   private case class StageState(numPartitions: Int) {
-    val authorizedCommitters = Array.fill[TaskIdentifier](numPartitions)(null)
+    val authorizedCommitters = Array.fill[CommitStatus](numPartitions)(null)
     val failures = mutable.Map[Int, mutable.Set[TaskIdentifier]]()
   }
 
@@ -100,6 +111,7 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
     val msg = AskPermissionToCommitOutput(stage, stageAttempt, partition, attemptNumber)
     coordinatorRef match {
       case Some(endpointRef) =>
+        endpointRef.send()
         ThreadUtils.awaitResult(endpointRef.ask[Boolean](msg),
           RpcUtils.askRpcTimeout(conf).duration)
       case None =>
@@ -109,6 +121,33 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
     }
   }
 
+  /**
+   * Called by tasks to update commit output success status.
+   *
+   * If a task attempt has been authorized to commit, after commit task success,
+   * task side should update commit status to true.
+   *
+   * @param stage the stage number
+   * @param partition the partition number
+   * @param attemptNumber how many times this task has been attempted
+   *                      (see [[TaskContext.attemptNumber()]])
+   * @return true if this task can update the commit output success status.
+   */
+  def commitSuccess(
+      stage: Int,
+      stageAttempt: Int,
+      partition: Int,
+      attemptNumber: Int): Unit = {
+    val msg = CommitOutputSuccess(stage, stageAttempt, partition, attemptNumber)
+    coordinatorRef match {
+      case Some(endpointRef) =>
+        endpointRef.send(msg)
+      case None =>
+        logError(
+          "commitSuccess called after coordinator was stopped (is SparkEnv shutdown in progress)?")
+    }
+  }
+
   /**
    * Called by the DAGScheduler when a stage starts. Initializes the stage's state if it hasn't
    * yet been initialized.
@@ -154,10 +193,17 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
         // Mark the attempt as failed to exclude from future commit protocol
         val taskId = TaskIdentifier(stageAttempt, attemptNumber)
         stageState.failures.getOrElseUpdate(partition, mutable.Set()) += taskId
-        if (stageState.authorizedCommitters(partition) == taskId) {
-          logDebug(s"Authorized committer (attemptNumber=$attemptNumber, stage=$stage, " +
-            s"partition=$partition) failed; clearing lock")
-          stageState.authorizedCommitters(partition) = null
+        val commitStatus = stageState.authorizedCommitters(partition)
+        if (commitStatus != null && commitStatus.taskIdent == taskId) {
+          if (commitStatus.status) {
+            throw new SparkException(s"Authorized committer (attemptNumber=$attemptNumber, " +
+              s"stage=$stage, partition=$partition) failed; but task commit success, " +
+              s"should fail the job")
+          } else {
+            logDebug(s"Authorized committer (attemptNumber=$attemptNumber, stage=$stage, " +
+              s"partition=$partition) failed; clearing lock")
+            stageState.authorizedCommitters(partition) = null
+          }
         }
     }
   }
@@ -186,11 +232,12 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
         if (existing == null) {
           logDebug(s"Commit allowed for stage=$stage.$stageAttempt, partition=$partition, " +
             s"task attempt $attemptNumber")
-          state.authorizedCommitters(partition) = TaskIdentifier(stageAttempt, attemptNumber)
+          state.authorizedCommitters(partition) =
+            CommitStatus(TaskIdentifier(stageAttempt, attemptNumber), false)
           true
         } else {
           logDebug(s"Commit denied for stage=$stage.$stageAttempt, partition=$partition: " +
-            s"already committed by $existing")
+            s"already committed by ${existing.taskIdent}")
           false
         }
       case None =>
@@ -200,6 +247,42 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
     }
   }
 
+  private[scheduler] def handleCommitOutputSuccess(
+      stage: Int,
+      stageAttempt: Int,
+      partition: Int,
+      attemptNumber: Int): Unit = synchronized {
+    stageStates.get(stage) match {
+      case Some(state) if attemptFailed(state, stageAttempt, partition, attemptNumber) =>
+        throw new SparkException(s"Authorized committer (attemptNumber=$attemptNumber, " +
+          s"stage=$stage, partition=$partition) failed; but task commit success, " +
+          s"should fail the job")
+      case Some(state) =>
+        val existing = state.authorizedCommitters(partition)
+        if (existing == null) {
+          throw new SparkException(s"Authorized committer (attemptNumber=$attemptNumber, " +
+          s"stage=$stage, partition=$partition) commit success; partition lock removed " +
+            s"before receive CommitOutputSuccess, should fail the job")
+        } else {
+          val taskIdent = existing.taskIdent
+          if (taskIdent.stageAttempt == stageAttempt && taskIdent.taskAttempt == attemptNumber) {
+            logDebug(s"Task Commit success for stage=$stage.$stageAttempt, " +
+              s"partition=$partition, task attempt $attemptNumber")
+            state.authorizedCommitters(partition) =
+              CommitStatus(TaskIdentifier(stageAttempt, attemptNumber), true)
+            true
+          } else {
+            throw new SparkException(s"Authorized committer (attemptNumber=$attemptNumber, " +
+              s"stage=$stage, partition=$partition) commit success; but partition lock not " +
+              s"consistent, should fail the job")
+          }
+        }
+      case None =>
+        logDebug(s"Commit update status failed for stage=$stage.$stageAttempt, " +
+          s"partition=$partition: stage already marked as completed.")
+    }
+  }
+
   private def attemptFailed(
       stageState: StageState,
       stageAttempt: Int,
@@ -223,6 +306,10 @@ private[spark] object OutputCommitCoordinator {
       case StopCoordinator =>
         logInfo("OutputCommitCoordinator stopped!")
         stop()
+
+      case CommitOutputSuccess(stage, stageAttempt, partition, attemptNumber) =>
+        outputCommitCoordinator.handleCommitOutputSuccess(stage, stageAttempt, partition,
+          attemptNumber)
     }
 
     override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {

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

@@ -270,6 +270,24 @@ class OutputCommitCoordinatorSuite extends SparkFunSuite with BeforeAndAfter {
       .stageStart(meq(retriedStage.head), any())
     verify(sc.env.outputCommitCoordinator).stageEnd(meq(retriedStage.head))
   }
+
+  test("SPARK-39195: Spark should use two step update of outputCommitCoordinator") {
+    var stage = 1
+    val taskAttempt = 1
+    val partition = 1
+
+    // Test receive CommitOutputSuccess but task failed.
+    outputCommitCoordinator.stageStart(stage, maxPartitionId = 1)
+    assert(outputCommitCoordinator.canCommit(stage, 1, partition, taskAttempt))
+    assert(!outputCommitCoordinator.canCommit(stage, 1, partition, taskAttempt + 1))
+    outputCommitCoordinator.commitSuccess(stage, 1, partition, taskAttempt)
+    val e1 = intercept[SparkException] {
+      outputCommitCoordinator.taskCompleted(stage, 1, partition, taskAttempt,
+        ExecutorLostFailure("0", exitCausedByApp = true, None))
+    }.getMessage
+    assert(e1.contains("Authorized committer (attemptNumber=1, stage=1, partition=1) failed; " +
+      "but task commit success, should fail the job"))
+  }
 }
 
 /**