[SPARK-3469] Call all TaskCompletionListeners even if some fail.

Note that this also changes the fault semantics of TaskCompletionListener.
Previously failures in TaskCompletionListeners would result in the task
being reported as failed. With this change, tasks won't be reported as failed
simply because the execution of TaskCompletionListener fails.

Author: rxin

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

@@ -77,6 +77,8 @@ class TaskContext(
   /**
    * 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.
+   * Exceptions in callbacks are however not thrown back upstream, i.e. tasks won't marked as
+   * failed even if completion callbacks fail to execute.
    *
    * An example use is for HadoopRDD to register a callback to close the input stream.
    */
@@ -104,7 +106,14 @@ class TaskContext(
   private[spark] def markTaskCompleted(): Unit = {
     completed = true
     // Process complete callbacks in the reverse order of registration
-    onCompleteCallbacks.reverse.foreach { _.onTaskCompletion(this) }
+    onCompleteCallbacks.reverse.foreach { listener =>
+      try {
+        listener.onTaskCompletion(this)
+      } catch {
+        case e: Throwable =>
+          logError("Error in TaskCompletionListener", e)
+      }
+    }
   }
 
   /** Marks the task for interruption, i.e. cancellation. */

core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala

@@ -26,7 +26,7 @@ import org.apache.spark.util.Utils
 
 class TaskContextSuite extends FunSuite with BeforeAndAfter with LocalSparkContext {
 
-  test("Calls executeOnCompleteCallbacks after failure") {
+  test("calls TaskCompletionListener after failure") {
     TaskContextSuite.completed = false
     sc = new SparkContext("local", "test")
     val rdd = new RDD[String](sc, List()) {
@@ -45,10 +45,21 @@ class TaskContextSuite extends FunSuite with BeforeAndAfter with LocalSparkConte
     }
     assert(TaskContextSuite.completed === true)
   }
+
+  test("all TaskCompletionListeners should be called even if some fail") {
+    val context = new TaskContext(0, 0, 0)
+    context.addTaskCompletionListener(_ => throw new Exception("blah"))
+    context.addTaskCompletionListener(_ => TaskContextSuite.callbackInvoked = true)
+    context.addTaskCompletionListener(_ => throw new Exception("blah"))
+    context.markTaskCompleted()
+    assert(TaskContextSuite.callbackInvoked === true)
+  }
 }
 
 private object TaskContextSuite {
   @volatile var completed = false
+
+  var callbackInvoked = false
 }
 
 private case class StubPartition(index: Int) extends Partition