address comments

Ngone51 · Ngone51 · commit c53961fd2d61 · 2019-02-26T21:46:25.000+08:00
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
@@ -858,12 +858,12 @@ private[spark] class TaskSchedulerImpl(
   private[scheduler] def markPartitionCompletedInAllTaskSets(
       stageId: Int,
       partitionId: Int,
-      taskInfo: Option[TaskInfo]) = {
+      taskInfo: TaskInfo) = {
     val finishedPartitions =
       stageIdToFinishedPartitions.getOrElseUpdate(stageId, new BitSet)
     finishedPartitions += partitionId
     taskSetsByStageIdAndAttempt.getOrElse(stageId, Map()).values.foreach { tsm =>
-      tsm.markPartitionCompleted(partitionId, taskInfo)
+      tsm.markPartitionCompleted(partitionId, Some(taskInfo))
     }
   }
 
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
@@ -196,9 +196,9 @@ private[spark] class TaskSetManager(
     // those tasks as finished here to avoid launching duplicate tasks, while
     // holding the TaskSchedulerImpl lock.
     // See SPARK-25250 and markPartitionCompletedInAllTaskSets()`
-    sched.stageIdToFinishedPartitions
-      .getOrElseUpdate(taskSet.stageId, new BitSet)
-      .foreach(markPartitionCompleted(_, None))
+    sched.stageIdToFinishedPartitions.get(taskSet.stageId).foreach {
+      finishedPartitions => finishedPartitions.foreach(markPartitionCompleted(_, None))
+    }
   }
 
   /**
@@ -798,7 +798,7 @@ private[spark] class TaskSetManager(
     }
     // There may be multiple tasksets for this stage -- we let all of them know that the partition
     // was completed.  This may result in some of the tasksets getting completed.
-    sched.markPartitionCompletedInAllTaskSets(stageId, tasks(index).partitionId, Some(info))
+    sched.markPartitionCompletedInAllTaskSets(stageId, tasks(index).partitionId, info)
     // This method is called by "TaskSchedulerImpl.handleSuccessfulTask" which holds the
     // "TaskSchedulerImpl" lock until exiting. To avoid the SPARK-7655 issue, we should not
     // "deserialize" the value when holding a lock to avoid blocking other threads. So we call

Original file line number	Diff line number	Diff line change
`@@ -858,12 +858,12 @@ private[spark] class TaskSchedulerImpl(`
`858`	`858`	`private[scheduler] def markPartitionCompletedInAllTaskSets(`
`859`	`859`	`stageId: Int,`
`860`	`860`	`partitionId: Int,`
`861`		`- taskInfo: Option[TaskInfo]) = {`
	`861`	`+ taskInfo: TaskInfo) = {`
`862`	`862`	`val finishedPartitions =`
`863`	`863`	`stageIdToFinishedPartitions.getOrElseUpdate(stageId, new BitSet)`
`864`	`864`	`finishedPartitions += partitionId`
`865`	`865`	`taskSetsByStageIdAndAttempt.getOrElse(stageId, Map()).values.foreach { tsm =>`
`866`		`- tsm.markPartitionCompleted(partitionId, taskInfo)`
	`866`	`+ tsm.markPartitionCompleted(partitionId, Some(taskInfo))`
`867`	`867`	`}`
`868`	`868`	`}`
`869`	`869`