[SPARK-3543] Write TaskContext in Java and expose it through a static accessor.

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

@@ -0,0 +1,176 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark;
+
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+import scala.Function0;
+import scala.Function1;
+import scala.Unit;
+import scala.collection.JavaConversions;
+
+import org.apache.spark.annotation.DeveloperApi;
+import org.apache.spark.executor.TaskMetrics;
+import org.apache.spark.util.TaskCompletionListener;
+import org.apache.spark.util.TaskCompletionListenerException;
+
+@DeveloperApi
+public class TaskContext implements Serializable {
+
+    public Integer stageId;
+    public Integer partitionId;
+    public Long attemptId;
+    public Boolean runningLocally;
+    public TaskMetrics taskMetrics;
+
+    public TaskContext(Integer stageId, Integer partitionId, Long attemptId, Boolean runningLocally,
+                       TaskMetrics taskMetrics) {
+        this.attemptId = attemptId;
+        this.partitionId = partitionId;
+        this.runningLocally = runningLocally;
+        this.stageId = stageId;
+        this.taskMetrics = taskMetrics;
+        taskContext.set(this);
+    }
+
+    public TaskContext(Integer stageId, Integer partitionId, Long attemptId,
+                       Boolean runningLocally) {
+        this.attemptId = attemptId;
+        this.partitionId = partitionId;
+        this.runningLocally = runningLocally;
+        this.stageId = stageId;
+        this.taskMetrics = TaskMetrics.empty();
+        taskContext.set(this);
+    }
+
+    public TaskContext(Integer stageId, Integer partitionId, Long attemptId) {
+        this.attemptId = attemptId;
+        this.partitionId = partitionId;
+        this.runningLocally = false;
+        this.stageId = stageId;
+        this.taskMetrics = TaskMetrics.empty();
+        taskContext.set(this);
+    }
+
+
+    private static ThreadLocal<TaskContext> taskContext =
+            new ThreadLocal<TaskContext>();
+
+    public static TaskContext get() {
+        return taskContext.get();
+    }
+
+    // List of callback functions to execute when the task completes.
+    private transient List<TaskCompletionListener> onCompleteCallbacks =
+            new ArrayList<TaskCompletionListener>();
+
+    // Whether the corresponding task has been killed.
+    private volatile Boolean interrupted = false;
+
+    // Whether the task has completed.
+    private volatile Boolean completed = false;
+
+    /**
+     * Checks whether the task has completed.
+     */
+    public Boolean isCompleted() {
+        return completed;
+    }
+
+    /**
+     * Checks whether the task has been killed.
+     */
+    public Boolean isInterrupted() {
+        return interrupted;
+    }
+
+
+    /**
+     * Add a (Java friendly) listener to be executed on task completion.
+     * This will be called in all situation - success, failure, or cancellation.
+     *
+     * An example use is for HadoopRDD to register a callback to close the input stream.
+     */
+    public TaskContext addTaskCompletionListener(TaskCompletionListener listener){
+        onCompleteCallbacks.add(listener);
+        return this;
+    }
+
+
+    /**
+     * Add a listener in the form of a Scala closure to be executed on task completion.
+     * This will be called in all situation - success, failure, or cancellation.
+     *
+     * An example use is for HadoopRDD to register a callback to close the input stream.
+     */
+    public TaskContext addTaskCompletionListener(final Function1<TaskContext, Unit> f) {
+        onCompleteCallbacks.add( new TaskCompletionListener() {
+
+            @Override
+            public void onTaskCompletion(TaskContext context) {
+                f.apply(context);
+            }
+        });
+        return this;
+    }
+
+    /**
+     * Add a callback function to be executed on task completion. An example use
+     * is for HadoopRDD to register a callback to close the input stream.
+     * Will be called in any situation - success, failure, or cancellation.
+     * @param f Callback function.
+     */
+    @Deprecated
+    public void addOnCompleteCallback(final Function0<Unit> f) {
+        onCompleteCallbacks.add(new TaskCompletionListener() {
+            @Override
+            public void onTaskCompletion(TaskContext context) {
+                f.apply();
+            }
+        });
+    }
+
+    /** Marks the task as completed and triggers the listeners. */
+    public void markTaskCompleted() throws TaskCompletionListenerException {
+        completed = true;
+        List<String> errorMsgs = new ArrayList<String>(2);
+        // Process complete callbacks in the reverse order of registration
+        List<TaskCompletionListener> revlist =
+                new ArrayList<TaskCompletionListener>(onCompleteCallbacks);
+        Collections.reverse(revlist);
+        for (TaskCompletionListener tcl : revlist){
+            try {
+                tcl.onTaskCompletion(this);
+            } catch (Throwable e){
+                errorMsgs.add(e.getMessage());
+            }
+        }
+
+        if (!errorMsgs.isEmpty()) {
+            throw new TaskCompletionListenerException(JavaConversions.asScalaBuffer(errorMsgs));
+        }
+    }
+
+    /** Marks the task for interruption, i.e. cancellation. */
+    public void markInterrupted() {
+        interrupted = true;
+    }
+}

core/src/main/scala/org/apache/spark/rdd/RDD.scala

@@ -619,6 +619,7 @@ abstract class RDD[T: ClassTag](
    * should be `false` unless this is a pair RDD and the input function doesn't modify the keys.
    */
   @DeveloperApi
+  @deprecated("use TaskContext.get", "1.2.0")
   def mapPartitionsWithContext[U: ClassTag](
       f: (TaskContext, Iterator[T]) => Iterator[U],
       preservesPartitioning: Boolean = false): RDD[U] = {

core/src/test/java/org/apache/spark/JavaAPISuite.java

@@ -776,7 +776,7 @@ public void persist() {
   @Test
   public void iterator() {
     JavaRDD<Integer> rdd = sc.parallelize(Arrays.asList(1, 2, 3, 4, 5), 2);
-    TaskContext context = new TaskContext(0, 0, 0, false, new TaskMetrics());
+    TaskContext context = new TaskContext(0, 0, 0L, false, new TaskMetrics());
     Assert.assertEquals(1, rdd.iterator(rdd.partitions().get(0), context).next().intValue());
   }
 

core/src/test/java/org/apache/spark/util/JavaTaskCompletionListenerImpl.java

@@ -30,10 +30,6 @@ public class JavaTaskCompletionListenerImpl implements TaskCompletionListener {
   public void onTaskCompletion(TaskContext context) {
     context.isCompleted();
     context.isInterrupted();
-    context.stageId();
-    context.partitionId();
-    context.runningLocally();
-    context.taskMetrics();
     context.addTaskCompletionListener(this);
   }
 }

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

@@ -94,7 +94,7 @@ class CacheManagerSuite extends FunSuite with BeforeAndAfter with EasyMockSugar
     }
 
     whenExecuting(blockManager) {
-      val context = new TaskContext(0, 0, 0, runningLocally = true)
+      val context = new TaskContext(0, 0, 0, true)
       val value = cacheManager.getOrCompute(rdd, split, context, StorageLevel.MEMORY_ONLY)
       assert(value.toList === List(1, 2, 3, 4))
     }