[SPARK-19263] DAGScheduler should avoid sending conflicting task set.

Author: jinxing

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

@@ -1193,7 +1193,14 @@ class DAGScheduler(
             }
 
             if (runningStages.contains(shuffleStage) && shuffleStage.pendingPartitions.isEmpty) {
-              markStageAsFinished(shuffleStage)
+              val noActiveTaskSetManager =
+                taskScheduler.rootPool == null ||
+                  !taskScheduler.rootPool.getSortedTaskSetQueue.exists {
+                    tsm => tsm.stageId == stageId && !tsm.isZombie
+                  }
+              if (shuffleStage.isAvailable || noActiveTaskSetManager) {
+                markStageAsFinished(shuffleStage)
+              }
               logInfo("looking for newly runnable stages")
               logInfo("running: " + runningStages)
               logInfo("waiting: " + waitingStages)
@@ -1218,7 +1225,9 @@ class DAGScheduler(
                 logInfo("Resubmitting " + shuffleStage + " (" + shuffleStage.name +
                   ") because some of its tasks had failed: " +
                   shuffleStage.findMissingPartitions().mkString(", "))
-                submitStage(shuffleStage)
+                if (noActiveTaskSetManager) {
+                  submitStage(shuffleStage)
+                }
               } else {
                 // Mark any map-stage jobs waiting on this stage as finished
                 if (shuffleStage.mapStageJobs.nonEmpty) {

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

@@ -38,7 +38,7 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.util.{CallSite, ThreadUtils, Utils}
 
 /**
- * Tests for the  entire scheduler code -- DAGScheduler, TaskSchedulerImpl, TaskSets,
+ * Tests for the entire scheduler code -- DAGScheduler, TaskSchedulerImpl, TaskSets,
  * TaskSetManagers.
  *
  * Test cases are configured by providing a set of jobs to submit, and then simulating interaction
@@ -648,4 +648,69 @@ class BasicSchedulerIntegrationSuite extends SchedulerIntegrationSuite[SingleCor
     }
     assertDataStructuresEmpty(noFailure = false)
   }
+
+  testScheduler("[SPARK-19263] DAGScheduler shouldn't resubmit active taskSet.") {
+    val a = new MockRDD(sc, 2, Nil)
+    val b = shuffle(2, a)
+    val shuffleId = b.shuffleDeps.head.shuffleId
+
+    def runBackend(): Unit = {
+      val (taskDescription, task) = backend.beginTask()
+      task.stageId match {
+        // ShuffleMapTask
+        case 0 =>
+          val stageAttempt = task.stageAttemptId
+          val partitionId = task.partitionId
+          (stageAttempt, partitionId) match {
+            case (0, 0) =>
+              val fetchFailed = FetchFailed(
+                DAGSchedulerSuite.makeBlockManagerId("hostA"), shuffleId, 0, 0, "ignored")
+              backend.taskFailed(taskDescription, fetchFailed)
+            case (0, 1) =>
+              // Wait until stage resubmission caused by FetchFailed is finished.
+              waitUntilConditionBecomeTrue(taskScheduler.runningTaskSets.size==2, 5000,
+                "Wait until stage is resubmitted caused by fetch failed")
+
+              // Task(stageAttempt=0, partition=1) will be bogus, because both two
+              // tasks(stageAttempt=0, partition=0, 1) run on hostA.
+              // Pending partitions are (0, 1) after stage resubmission,
+              // then change to be 0 after this bogus task.
+              backend.taskSuccess(taskDescription, DAGSchedulerSuite.makeMapStatus("hostA", 2))
+            case (1, 1) =>
+              // Wait long enough until Success of task(stageAttempt=1 and partition=0)
+              // is handled by DAGScheduler.
+              Thread.sleep(5000)
+              // Task(stageAttempt=1 and partition=0) will cause stage resubmission,
+              // because shuffleStage.pendingPartitions.isEmpty,
+              // but shuffleStage.isAvailable is false.
+              backend.taskSuccess(taskDescription, DAGSchedulerSuite.makeMapStatus("hostB", 2))
+            case _ =>
+              backend.taskSuccess(taskDescription, DAGSchedulerSuite.makeMapStatus("hostB", 2))
+          }
+        // ResultTask
+        case 1 => backend.taskSuccess(taskDescription, 10)
+      }
+    }
+
+    withBackend(runBackend _) {
+      val jobFuture = submit(b, (0 until 2).toArray)
+      val duration = Duration(15, SECONDS)
+      awaitJobTermination(jobFuture, duration)
+    }
+    assert(results === (0 until 2).map { _ -> 10}.toMap)
+  }
+
+  def waitUntilConditionBecomeTrue(condition: => Boolean, timeout: Long, msg: String): Unit = {
+    val finishTime = System.currentTimeMillis() + timeout
+    while (System.currentTimeMillis() < finishTime) {
+      if (condition) {
+        return
+      }
+      // Sleep rather than using wait/notify, because this is used only for testing and wait/notify
+      // add overhead in the general case.
+      Thread.sleep(10)
+    }
+    throw new TimeoutException(
+      s"Condition '$msg' failed to become true before $timeout milliseconds elapsed")
+  }
 }