Fixed the hanging in JobCancellationSuite.

apache · Aug 20, 2014 · c414c36 · c414c36
1 parent b3e2eed
commit c414c36
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
@@ -677,8 +677,10 @@ class DAGScheduler(
   }
 
   private[scheduler] def handleBeginEvent(task: Task[_], taskInfo: TaskInfo) {
-    val stageInfo = stageIdToStage(task.stageId).latestInfo
-    listenerBus.post(SparkListenerTaskStart(task.stageId, stageInfo.attemptId, taskInfo))
+    // Note that there is a chance that this task is launched after the stage is cancelled.
+    // In that case, we wouldn't have the stage anymore in stageIdToStage.
+    val stageAttemptId = stageIdToStage.get(task.stageId).map(_.latestInfo.attemptId).getOrElse(-1)
+    listenerBus.post(SparkListenerTaskStart(task.stageId, stageAttemptId, taskInfo))
     submitWaitingStages()
   }