Move registration of task with BlockManager into Task.run()

Author: JoshRosen

core/src/main/scala/org/apache/spark/executor/Executor.scala

@@ -207,8 +207,6 @@ private[spark] class Executor(
         logDebug("Task " + taskId + "'s epoch is " + task.epoch)
         env.mapOutputTracker.updateEpoch(task.epoch)
 
-        env.blockManager.registerTask(taskId)
-
         // Run the actual task and measure its runtime.
         taskStart = System.currentTimeMillis()
         var threwException = true

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

@@ -64,6 +64,7 @@ private[spark] abstract class Task[T](
       taskAttemptId: Long,
       attemptNumber: Int,
       metricsSystem: MetricsSystem): T = {
+    SparkEnv.get.blockManager.registerTask(taskAttemptId)
     context = new TaskContextImpl(
       stageId,
       partitionId,

core/src/test/scala/org/apache/spark/CacheManagerSuite.scala

@@ -87,6 +87,7 @@ class CacheManagerSuite extends SparkFunSuite with LocalSparkContext with Before
     val context = TaskContext.empty()
     try {
       TaskContext.setTaskContext(context)
+      sc.env.blockManager.registerTask(0)
       cacheManager.getOrCompute(rdd3, split, context, StorageLevel.MEMORY_ONLY)
       assert(context.taskMetrics.updatedBlockStatuses.size === 2)
     } finally {