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


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

package spark

import scala.collection.mutable.ArrayBuffer


class TaskContext(val stageId: Int, val splitId: Int, val attemptId: Long) extends Serializable {

  @transient val onCompleteCallbacks = new ArrayBuffer[() => Unit]

  // 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.
  def addOnCompleteCallback(f: () => Unit) {
    onCompleteCallbacks += f
  }

  def executeOnCompleteCallbacks() {
    onCompleteCallbacks.foreach{_()}
  }
}