Moved setting the thread local to where TaskContext is instantiated.

Author: ScrapCodes

core/src/main/java/org/apache/spark/TaskContext.java

@@ -39,10 +39,10 @@
 @DeveloperApi
 public class TaskContext implements Serializable {
 
-  private Integer stageId;
-  private Integer partitionId;
-  private Long attemptId;
-  private Boolean runningLocally;
+  private int stageId;
+  private int partitionId;
+  private long attemptId;
+  private boolean runningLocally;
   private TaskMetrics taskMetrics;
 
   /**
@@ -63,7 +63,6 @@ public TaskContext(Integer stageId, Integer partitionId, Long attemptId, Boolean
     this.runningLocally = runningLocally;
     this.stageId = stageId;
     this.taskMetrics = taskMetrics;
-    taskContext.set(this);
   }
 
 
@@ -84,7 +83,6 @@ public TaskContext(Integer stageId, Integer partitionId, Long attemptId,
     this.runningLocally = runningLocally;
     this.stageId = stageId;
     this.taskMetrics = TaskMetrics.empty();
-    taskContext.set(this);
   }
 
 
@@ -103,12 +101,15 @@ public TaskContext(Integer stageId, Integer partitionId, Long attemptId) {
     this.runningLocally = false;
     this.stageId = stageId;
     this.taskMetrics = TaskMetrics.empty();
-    taskContext.set(this);
   }
 
   private static ThreadLocal<TaskContext> taskContext =
     new ThreadLocal<TaskContext>();
 
+  public static void setTaskContext(TaskContext tc) {
+    taskContext.set(tc);
+  }
+
   public static TaskContext get() {
     return taskContext.get();
   }
@@ -212,19 +213,19 @@ public void markInterrupted() {
     interrupted = true;
   }
 
-  public Integer stageId() {
+  public int stageId() {
     return stageId;
   }
 
-  public Integer partitionId() {
+  public int partitionId() {
     return partitionId;
   }
 
-  public Long attemptId() {
+  public long attemptId() {
     return attemptId;
   }
 
-  public Boolean runningLocally() {
+  public boolean runningLocally() {
     return runningLocally;
   }
 

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

@@ -634,7 +634,8 @@ class DAGScheduler(
       val rdd = job.finalStage.rdd
       val split = rdd.partitions(job.partitions(0))
       val taskContext =
-        new TaskContext(job.finalStage.id, job.partitions(0), 0, runningLocally = true)
+        new TaskContext(job.finalStage.id, job.partitions(0), 0, true)
+      TaskContext.setTaskContext(taskContext)
       try {
         val result = job.func(taskContext, rdd.iterator(split, taskContext))
         job.listener.taskSucceeded(0, result)

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

@@ -45,7 +45,8 @@ import org.apache.spark.util.Utils
 private[spark] abstract class Task[T](val stageId: Int, var partitionId: Int) extends Serializable {
 
   final def run(attemptId: Long): T = {
-    context = new TaskContext(stageId, partitionId, attemptId, runningLocally = false)
+    context = new TaskContext(stageId, partitionId, attemptId, false)
+    TaskContext.setTaskContext(context)
     context.taskMetrics.hostname = Utils.localHostName()
     taskThread = Thread.currentThread()
     if (_killed) {