SparkCore — Task执行源码分析之Task.run()源码分析

Task.run()源码分析

  上一篇博客分析了TaskRunner.run()源码，它里面有两个比较重要的方法一个是Task.run() — task的执行，还有就是task执行结束后将执行结果发送给Driver的StatusUpdate()，这里我们来分析Task.run()方法：

final def run(
    taskAttemptId: Long,
    attemptNumber: Int,
    metricsSystem: MetricsSystem)
  : (T, AccumulatorUpdates) = {
    // 创建TaskContext，就是task的执行上下文
    // 里面包含执行task所需的一些全局性变量，比如task重试次数，task的内存管理，task的stage
    // task要处理的是RDD的哪个partition等
    context = new TaskContextImpl(
      stageId,
      partitionId,
      taskAttemptId,
      attemptNumber,
      taskMemoryManager,
      metricsSystem,
      internalAccumulators,
      runningLocally = false)
    TaskContext.setTaskContext(context)
    context.taskMetrics.setHostname(Utils.localHostName())
    context.taskMetrics.setAccumulatorsUpdater(context.collectInternalAccumulators)
    taskThread = Thread.currentThread()
    if (_killed) {
      kill(interruptThread = false)
    }
    try {
      // 调用抽象方法 runTask()
      // 调用抽象方法，意味着，这个类仅仅是一个模板类，或者抽象父类
      // 仅仅封装了一些子类通用的数据和操作，关键的操作全部依赖于子类的实现
      // Task的子类有两种，ShuffleMapTask和ResultTask
      (runTask(context), context.collectAccumulators())
    } finally {
      context.markTaskCompleted()
      try {
        Utils.tryLogNonFatalError {
          // 获取BlockManager的内存管理器，释放内存
          SparkEnv.get.blockManager.memoryStore.releaseUnrollMemoryForThisTask()
          val memoryManager = SparkEnv.get.memoryManager
          memoryManager.synchronized { memoryManager.notifyAll() }
        }
      } finally {
        TaskContext.unset()
      }
    }
  }

  从源码中看，首先创建执行task的上下文，里面封装了一些执行task所需的一些全局变量，比如task待处理的RDD的partiton ID，task重试次数，task的内存管理器等。接着调用抽象方法runTask()，由于Task类是抽象类，因此真正执行的是在子类中，它只封装了子类通用的一些处理逻辑和数据，具体的执行还是子类来执行，Task的子类有两种：ShuffleMapTask和ResultTask。
  我们先看ShuffleMapTask中的runTask()方法：

 override def runTask(context: TaskContext): MapStatus = {
    // Deserialize the RDD using the broadcast variable.
    val deserializeStartTime = System.currentTimeMillis()
    // 对Task要处理的RDD相关的数据，做一些反序列化操作
    val ser = SparkEnv.get.closureSerializer.newInstance()
    // RDD是怎么拿到的？多个task运行在多个executor中，都是并行运行，或者并发运行的
    // 但是一个stage的task要处理的RDD是一样的，所以task怎么拿到自己负责的要处理的RDD的数据？
    // 会通过broadcast variable，具体源码在DAGScheduler中的submitMissingTasks中，下面给出这个代码的片段，大家可以看一下。
    val (rdd, dep) = ser.deserialize[(RDD[_], ShuffleDependency[_, _, _])](
      ByteBuffer.wrap(taskBinary.value), Thread.currentThread.getContextClassLoader)
    _executorDeserializeTime = System.currentTimeMillis() - deserializeStartTime

    metrics = Some(context.taskMetrics)
    var writer: ShuffleWriter[Any, Any] = null
    try {
      // 获取shufflemanager
      val manager = SparkEnv.get.shuffleManager
      // 获取ShuffleManager的ShuffleWriter
      writer = manager.getWriter[Any, Any](dep.shuffleHandle, partitionId, context)
      // 首先调用的RDD的iterator()方法，并且传入了，当前task要处理的哪个partition
      // 所以核心逻辑就在RDD的iterator()方法这里，实现了针对RDD的某个partition，执行我们自己定义的算子或函数
      writer.write(rdd.iterator(partition, context).asInstanceOf[Iterator[_ <: Product2[Any, Any]]])
      // 返回结果，MapStatus，里面封装了ShuffleMapTask计算后的数据存储在哪里，其实就是BlockManager相关的信息
      // BlockManager是Spark底层的内存、数据、磁盘数据管理的组件
      writer.stop(success = true).get
    } catch {
      case e: Exception =>
        // 省略异常处理代码
    }
  }

  首先ShuffleMapTask的runTask()方法有返回值，它的返回值是MapStatus。接着会调用ShuffleWriter的Write方法，将Task算子运行结果写入map output文件中去，在分析Shuffle流程的时候会讲到。write()方法中调用了RDD的iterator方法，这个方法里面就会执行我们定义的算子函数。它传入了当前task要处理哪个partition。这里就是runTask的核心逻辑。
  下面是RDD的iterator方法：

final def iterator(split: Partition, context: TaskContext): Iterator[T] = {
    // 如果有本地化级别
    if (storageLevel != StorageLevel.NONE) {
      // CacheManager是从持久化的RDD中读取当前计算RDD需要的数据
      SparkEnv.get.cacheManager.getOrCompute(this, split, context, storageLevel)
    } else {
      // 假如没有进行cache，那么进行RDD的partition的计算
      computeOrReadCheckpoint(split, context)
    }
  }

  首先看这个RDD是否被cache了，假如没有的话，就执行对RDD的partition的计算，其中关于CacheManager后面再分析，其实从这里就能看出，假如对反复使用的RDD进行cache或者checkpoint的话，那么就不需要重复计算，这是性能优化的地方，后面再讲。

private[spark] def computeOrReadCheckpoint(split: Partition, context: TaskContext): Iterator[T] =
  {
    // 是否是checkpoint
    if (isCheckpointedAndMaterialized) {
      firstParent[T].iterator(split, context)
    } else {
      // 计算，它是抽象方法。
      compute(split, context)
    }
  }

  这里首先也是判断是否checkpoint了，这里也不进行分析，下面再进行分析，假如没有checkpoint，那么就直接compute()计算。这里compute()是抽象方法，我们看以它的一个常见的子类MapPartitonsRDD的compute进行分析。

override def compute(split: Partition, context: TaskContext): Iterator[U] =
    // 这里的f，就是我们自己定义的算子和函数，Spark进行了一些封装，还实现了一些其他的逻辑
    // 调用到这里其实就是针对RDD的partition，执行自定义的计算操作，并返回新的RDD的partition的数据
    f(context, split.index, firstParent[T].iterator(split, context))

  这里就很有含义了，compute就是针对RDD中的某个partition执行我们给这个RDD定义的算子和函数。这个方法里面的f，可以理解为我们自己定义的算子和函数，这里会返回新的RDD的partition的数据。
  以上就是我们Task.run()方法的执行逻辑，总结一下，主要是调用ShuffleManager的ShuffleWriter的write方法，传入RDD的iterator()方法，里面执行我们定义的算子和函数，然后将执行结果写入map output文件中，然后获取执行的状态MapStatus，并返回。