Spark源码精读一一心跳机制

注册端点

    // We need to register "HeartbeatReceiver" before "createTaskScheduler" because Executor will
    // retrieve "HeartbeatReceiver" in the constructor. (SPARK-6640)
    // 创建心跳端点并注册到Dispatcher消息派发器中
    _heartbeatReceiver = env.rpcEnv.setupEndpoint(
      HeartbeatReceiver.ENDPOINT_NAME, new HeartbeatReceiver(this))

这里要理解一个RpcEndpoint和NettyRpcEndpointRef和EndpointData概念及其意义。
NettyRpcEndpointRef： 中间对象，RpcEndpoint端点与Dispatcher消息派发器之间的桥梁。不管Spark中定义了多少个Endpoint，注册端点时，都要通过NettyRpcEndpointRef交由Dispatcher派发处理。
RpcEndpoint： 端点的具体实现，比如，实现发送具体的消息。
EndpointData： RpcEndpoint端点与NettyRpcEndpointRef端点的关系对象，以便MessageLoop处理消息时，知道是哪一个端点的消息。

向Dispatcher注册端点，其实就是向Dispatcher的端点队列中添加该端点。

  // Track the receivers whose inboxes may contain messages.
  // 这是一个线程安全的共享阻塞队列
  private val receivers = new LinkedBlockingQueue[EndpointData]

当main线程将HearbeatReceiver端点注册到Dispatcher中（通过该队列的offer操作添加到Dispatcher的共享阻塞队列receivers）会唤醒所有阻塞在take操作的派发线程从队列中获取端点。由于take操作内部也是ReentrantLock，因此，所有的线程中只有一个能获取该锁，即有且只有一个线程能够获取到，通过注册的端点。这一过程详见LinkedBlockingQueue。

当获取到心跳端点的EndpointData后，通过其inbox来处理其启动消息。该启动消息是在创建EndpointData时添加到Inbox的消息集合中去。

  ...
  // OnStart should be the first message to process
  inbox.synchronized {
    messages.add(OnStart)
  }
  ...
  ...
  data.inbox.process(Dispatcher.this)

从该端点所属的消息集合messages中，获取消息，此时是OnStart消息。经过消息匹配逻辑，调用心跳端点的启动方法endpoint.onStart()。

启动定时检查过期主机

 override def onStart(): Unit = {
    timeoutCheckingTask = eventLoopThread.scheduleAtFixedRate(new Runnable {
      override def run(): Unit = Utils.tryLogNonFatalError {
        // 隔checkTimeoutIntervalMs微妙，询问过期断开的主机消息
        Option(self).foreach(_.ask[Boolean](ExpireDeadHosts))
      }
    }, 0, checkTimeoutIntervalMs, TimeUnit.MILLISECONDS)
  }

通过对OnStart和ExpireDeadHosts消息观察，发现，Scala中的消息定义即case object或者case class，比如消息RpcEndpointVerifier.CheckExistence(endpointRef.name)。 这两种消息的不同在于，case class可以带参数，而case object不带参数。

private[spark] case object TaskSchedulerIsSet
private[spark] case object ExpireDeadHosts
// 消息
//org.apache.spark.rpc.netty.RpcEndpointVerifier case class CheckExistence(name: String) extends Product with Serializable with Product with Serializable with Object
/** A message used to ask the remote [[RpcEndpointVerifier]] if an `RpcEndpoint` exists. */
RpcEndpointVerifier.CheckExistence(endpointRef.name)

询问TaskScheduler已经创建

当TaskScheduler对象创建完后，心跳端点会询问TaskScheduler是否已创建。

// _heartbeatReceiver其实是NettyRpcEndpointRef。
_heartbeatReceiver.ask[Boolean](TaskSchedulerIsSet)

通过ask操作，来探究Spark的消息发送到响应的处理细节。

 /**
   * Send a message to the corresponding [[RpcEndpoint.receiveAndReply)]] and return a [[Future]] to
   * receive the reply within a default timeout.
   *
   * This method only sends the message once and never retries.
   */
  def ask[T: ClassTag](message: Any): Future[T] = 
  		// 调用的其实是NettyRpcEndpointRef的。因为是经过它将消息交给Dispatcher的
  		ask(message, defaultAskTimeout)

继续debug，程序走到NettyRpcEndpointRef的ask函数

  override def ask[T: ClassTag](message: Any, timeout: RpcTimeout): Future[T] = {
    // 将地址跟发送的消息封装成RequestMessage，并由nettyRpcEnv的ask函数将消息发送。
    nettyEnv.ask(RequestMessage(nettyEnv.address, this, message), timeout)
  }

由此可见，Spark是通过NettyRpcEndpointRef将RpcEndpoint与Netty进行了关联。
由于Spark下的NettyRpc的ask函数十分复杂。将代码拆分一步一步分析。

    // 许可对象，类似Netty的Promise机制。
	val promise = Promise[Any]()
	// 这个发送到心跳端点的地址 message.receiver = NettyRpcEndpointRef。
	// 发送地址：是NettyRpcEnv处的地址，为nettyEnv.address
    val remoteAddr = message.receiver.address
	// 响应失败的函数
    def onFailure(e: Throwable): Unit = {
      if (!promise.tryFailure(e)) {
        logWarning(s"Ignored failure: $e")
      }
    }
	// 响应成功的函数
    def onSuccess(reply: Any): Unit = reply match {
      case RpcFailure(e) => onFailure(e)
      case rpcReply =>
        if (!promise.trySuccess(rpcReply)) {
          logWarning(s"Ignored message: $reply")
        }
    }

假定，当前分析都是Spark运行在本地模式下。
通过远程地址跟当前地址比较remoteAddr == address，相等，则是本地调用，而不是远程调用。

	// 重新定义一个许可
	// 这个Promise其实是一个Object，定义了apply方法
	// def apply[T](): Promise[T] = new impl.Promise.DefaultPromise[T]()
	val p = Promise[Any]()
	// p.future的意义在于确定调用函数onComplete的类型为DefaultPromise
	// def future: this.type = this
    p.future.onComplete {
      // 许可p下的回调函数体
      case Success(response) => onSuccess(response)
      case Failure(e) => onFailure(e)
    }(ThreadUtils.sameThread) // 隐式参数显示调用

onComplete函数在DefaultPromise定义

	def onComplete[U](func: Try[T] => U)(implicit executor: ExecutionContext): Unit = {
      val preparedEC = executor.prepare()
      // 将执行器和函数变量封装为回调对象
      val runnable = new CallbackRunnable[T](preparedEC, func)
      // 派发或者添加回调对象
      dispatchOrAddCallback(runnable)
    }
    
    @tailrec
    private def dispatchOrAddCallback(runnable: CallbackRunnable[T]): Unit = {
      // 当第一次进入时，getState方法获取的是空集合，与第三个匹配
      getState match {
        // 执行onComplete函数中的函数表达式
        case r: Try[_]          => runnable.executeWithValue(r.asInstanceOf[Try[T]])
        case _: DefaultPromise[_] => compressedRoot().dispatchOrAddCallback(runnable)
        case listeners: List[_] => if (updateState(listeners, runnable :: listeners)) () else dispatchOrAddCallback(runnable)
      }
    }

接下来，派发器将本消息TaskSchedulerIsSet和许可进行本地投递。

	// 派发器将消息message和许可p进行投递
    dispatcher.postLocalMessage(message, p)

通过本地端点投递消息

  /** Posts a message sent by a local endpoint. */
  def postLocalMessage(message: RequestMessage, p: Promise[Any]): Unit = {
    val rpcCallContext =
      // 创建本地的NettyRpc调用上下文
      new LocalNettyRpcCallContext(message.senderAddress, p)
    // 发送者的消息封装
    val rpcMessage = RpcMessage(message.senderAddress, message.content, rpcCallContext)
    // 向接受者投递发送者发送的message.content=TaskSchedulerIsSet消息
    postMessage(message.receiver.name, rpcMessage, (e) => p.tryFailure(e))
  }
  
  /**
   * Posts a message to a specific endpoint.
   * 投递消息到指定的端点
   * @param endpointName name of the endpoint.
   * @param message the message to post
   * @param callbackIfStopped callback function if the endpoint is stopped.
   */
  private def postMessage(
      endpointName: String,
      message: InboxMessage,
      callbackIfStopped: (Exception) => Unit): Unit = {
    val error = synchronized {
      val data = endpoints.get(endpointName)
      if (stopped) {
        Some(new RpcEnvStoppedException())
      } else if (data == null) {
        Some(new SparkException(s"Could not find $endpointName."))
      } else {
        // 将消息message投递到接收者端点中
        // 此时的data.inbox实际是接收端点的
        data.inbox.post(message)
        // 将端点添加到端点接受者队列中
        receivers.offer(data)
        None
      }
    }
    // We don't need to call `onStop` in the `synchronized` block
    error.foreach(callbackIfStopped)
  }

当消息派发器Dispatcher的端点队列重新有值，即上步向接收端点投递消息时，向队列添加了接收者的端点。进行接收者端点的inbox的消息process处理。

	case RpcMessage(_sender, content, context) =>
      try {
      	// 此时的endpoint是接收端点的endpoint = HeartbeatReceiver
        endpoint.receiveAndReply(context).applyOrElse[Any, Unit](content, { msg =>
          throw new SparkException(s"Unsupported message $message from ${_sender}")
        })
      } catch {
        case NonFatal(e) =>
          context.sendFailure(e)
          // Throw the exception -- this exception will be caught by the safelyCall function.
          // The endpoint's onError function will be called.
          throw e
      }

HeartbeatReceiver接收消息并回复receiveAndReply：

	case TaskSchedulerIsSet =>
      scheduler = sc.taskScheduler
      // 此时的context其实就是RpcMessage消息投递时的NettyRpcCallContext。
      // 即LocalNettyRpcCallContext
      context.reply(true)

LocalNettyRpcCallContext的响应恢复：

  override protected def send(message: Any): Unit = {
    p.success(message)
  }

Promise的success函数：

 /** Completes the promise with a value.
   *
   *  @param value The value to complete the promise with.
   *
   *  $promiseCompletion
   */
  def success(@deprecatedName('v) value: T): this.type = complete(Success(value))

 /** Completes the promise with either an exception or a value.
   *
   *  @param result     Either the value or the exception to complete the promise with.
   *  
   *  $promiseCompletion
   */
  def complete(result: Try[T]): this.type =
    if (tryComplete(result)) this else throw new IllegalStateException("Promise already completed.")

DefaultPromise下的tryComplete：

 def tryComplete(value: Try[T]): Boolean = {
      val resolved = resolveTry(value)
      tryCompleteAndGetListeners(resolved) match {
        case null             => false
        case rs if rs.isEmpty => true
        // 执行回调 此处的r是之前创建的CallbackRunnable
        case rs               => rs.foreach(r => r.executeWithValue(resolved)); true
      }
    }
    
 def executeWithValue(v: Try[T]): Unit = {
    require(value eq null) // can't complete it twice
    value = v
    // Note that we cannot prepare the ExecutionContext at this point, since we might
    // already be running on a different thread!
    // 触发执行回调的onComplete的函数体表达式内容
    // 这里要注意一下：
    // 这里回调的是许可p的函数onComplete的函数表达式
    // 因为消息投递时，投递的许可就是p而不是promise。
    // 所以执行到这，实际就是执行消息投递之前的许可p的onComplete函数的函数表达式
    // p.future.onComplete {
    //  // 许可p下的回调函数体，由接收者的响应reply(context)消息时触发
    //  case Success(response) => onSuccess(response) 消息超时的许可逻辑
    //  case Failure(e) => onFailure(e)
    // }(ThreadUtils.sameThread) // 隐式参数显示调用
    try executor.execute(this) catch { case NonFatal(t) => executor reportFailure t }
  }