[以浪为码]Spark源码阅读02 - RPC模块

Spark 的 RPC 模块是建立在 network 模块之上，虽然 network 提供了远程调用与数据流传输，但是 RPC 提供了更加方便的编程方式与性能提升。本文通过阅读 RPC 模块的代码，来了解其实现。

在此之前，建议提前了解一下 network 模块。这里简单介绍一下，详细请看[以浪为码]Spark源码阅读01-网络传输 network。

网络传输模块实现了 RPC 、流数据传输与数据块传输，主要分为客户端与服务端，客户端TransportClient提供了相应的请求发送的方法，并且在请求时需要使用回调来设置响应处理；服务端 TransportServer 接收请求，处理请求并返回响应，对于请求的处理，模块用户需要实现 RpcHandler 来定义。两者都从 TransportContext 创建（客户端为客户端池）。

RpcEndpoint rpc终点

每一个 RPC 发送之后绕来绕去后都需要到达一个目的地或者终点，RpcEndpoint 接口就代表的是这个终点， Spark 中的 Master、Worker 都是一个RpcEndpoint。他负责定义给定消息到他这里之后所触发的操作。且每个 RpcEndpoint 在系统中都有自己的一个名字。

RpcEndpoint 的生命周期为 constructor -> onStart -> receive* -> onStop ，如果 RpcEndpoint 抛出了错误， RpcEndpoint 就会调用 onError 方法。

上面提到的接口方法代码如下：

code 1

  def onStart(): Unit = {
    // By default, do nothing.
  }

  def onStop(): Unit = {
    // By default, do nothing.
  }
  
  def receive: PartialFunction[Any, Unit] = {
    case _ => throw new SparkException(self + " does not implement 'receive'")
  }

  def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
    case _ => context.sendFailure(new SparkException(self + " won't reply anything"))
  }
  def onError(cause: Throwable): Unit = {
    // By default, throw e and let RpcEnv handle it
    throw cause
  }

由代码可知，receve* 方法返回一个 PartialFunction[Any, Unit]，即 Scala 中的偏函数，可以方便的定义不同消息的处理方法。

RpcEndpointRef

RpcEndpointRef 抽象类是 RpcEndpoint 的引用，负责向 RpcEndpoint 发送消息，他知道 RpcEndpoint 地址与名字。这里列一下消息发送的方法：

code 2

// 发送一个不需要回复的消息。
def send(message: Any): Unit

// 发送一个消息，对应于 `RpcEndpoint.receiveAndReply` 方法，他返回一个 Future 在指定的超时时间内接收回复。
// 这个方法只发送一次消息，不会重试
def ask[T: ClassTag](message: Any, timeout: RpcTimeout): Future[T]

// 使用了默认超时时间的 `ask`,超时时间使用 spark.rpc.askTimeout 指定，默认 120s
def ask[T: ClassTag](message: Any): Future[T] = ask(message, defaultAskTimeout)

// `ask` 的同步实现，需要指定超时时间。
def askSync[T: ClassTag](message: Any, timeout: RpcTimeout): T = {
    val future = ask[T](message, timeout)
    timeout.awaitResult(future)
  }
  
// `ask` 的同步实现，使用默认超时时间。超时时间使用 spark.rpc.askTimeout 指定，默认 120s
def askSync[T: ClassTag](message: Any): T = askSync(message, defaultAskTimeout)

RpcEnv 与 NettyRpcEnv

接下来介绍 RPC 运行环境的抽象类 RpcEnv 与他在 Spark 中的唯一实现NettyRpcEnv。RpcEndpoint 带一个名字注册到RpcEnv上以用于接收数据。RpcEnv 会处理来自 RpcEndpointRef 与 远程节点的消息，并将消息分发给响应的RpcEndpoint 。对于 RpcEnv 抛出的未捕获的异常，RpcEnv 会使用 RpcCallContext.sendFailure 将异常报告给消息发送者，或者在没有发送者或 NotSerializableException 的时候打印出日志。

我们直接看实现 NettyRpcEnv。

NettyRpcEnv 的成员有一个 TransportContext 传输上下文，当然还有从 TransportContext 中得到的传输服务端 TransportServer 与 传输客户端工厂TransportClientFactory（用于 RPC），此外还有一个是用于专门的文件下载的客户端工厂，他是为了避免通信阻塞，他与用于RPC的客户端工厂的不同仅仅是创建他们的TransportContext 配置不同。

RPC消息发送

RpcEndpointRef 的实现 NettyRpcEndpointRef 发送消息的相关代码如下，可见他是直接调用NettyRpcEnv 的对应方法：

code 3

private[netty] class NettyRpcEndpointRef(
    @transient private val conf: SparkConf,
    private val endpointAddress: RpcEndpointAddress,
    @transient @volatile private var nettyEnv: NettyRpcEnv) extends RpcEndpointRef(conf)
   ...

  override def ask[T: ClassTag](message: Any, timeout: RpcTimeout): Future[T] = {
    nettyEnv.ask(new RequestMessage(nettyEnv.address, this, message), timeout)
  }
  override def send(message: Any): Unit = {
    require(message != null, "Message is null")
    nettyEnv.send(new RequestMessage(nettyEnv.address, this, message))
  }
  ...
}

Outbox 发件箱

在 NettyRpcEnv 中，管理着叫做发件箱Outbox的组件,如下所示，一个RPC的具体地址 RpcAddress 对应一个 Outbox。

code 4

  private val outboxes = new ConcurrentHashMap[RpcAddress, Outbox]()

Outbox 是一个非阻塞批量发送消息的发送器，NettyRpcEndpointRef 发送到远程地址的消息都会由Outbox 进行实际发送。Outbox 维护着TransportClient用于发送消息，一个消息队列用于存储消息， Outbox 接受到消息后，消息主要会经历两个阶段,如 code 5 所示：

1. 消息先放到消息队列中。
2. 将消息从消息队列中倾倒（发送）出去。倾倒发送的代码这里不贴了，这里描述一下，Outbox会先检查当前是否有其他线程在倾倒消息，如果没有则从队列 poll 消息，串行将消息发送出去，直到队列为空才。可以想象在 Outbox 正由某个线程倾倒消息的时候，其他线程可以正常的往 Outbox 中发送消息，而不会阻塞。这种使用发送数据的线程同时批量发送数据的设计，而不使用维护一个轮询线程来实现批量发送，即实现了批量发送，也不至于Outbox 成为一个服务式对象，符合消息发送后即止的普遍语义。

code 5

private[netty] class Outbox(nettyEnv: NettyRpcEnv, val address: RpcAddress) {

  private val messages = new java.util.LinkedList[OutboxMessage]
  ...
  def send(message: OutboxMessage): Unit = {
    val dropped = synchronized {
      if (stopped) {
        true
      } else {
        // 放入消息列表
        messages.add(message)
        false
      }
    }
    if (dropped) {
      message.onFailure(new SparkException("Message is dropped because Outbox is stopped"))
    } else {
      drainOutbox()
    }
  }
  ...
}

当然以上的发送过程对 NettyRpcEndpointRef 来说是透明的。

RPC 消息接收

NettyRpcEnv 里比较重要的成员是 Dispatcher 分发器， 是他负责将发送到 NettyRpcEnv 的消息分发到对应的 EndPoint 上，分发器分发的消息包括远程发送来的消息以及本地RpcEndpointRef消息。

Dispatcher 使用了一个称为收件箱Inbox设计，以实现消息的批量发送。一个 Inbox 负责一个 Endpoint 的消息， Dispatcher 使用一个内部类 EndpointData 维护他们之间的关系, 如 code 6 所示， 没将当 Dispatcher 收到一个消息，就根据name将消息放入对应的 EndpointData 的 inbox 中的，并将 EndpointData （引用）放入一个阻塞队列，如 code 6 postMessage 方法所示。

很显然，阻塞队列就是需要被消费的，与此同时，Dispatcher 会使用spark.rpc.netty.dispatcher.numThreads个线程去消费阻塞队列 receivers， 触发其 EndpointData 中 Inbox 的消息处理方法，见 code 6 的 MessageLoop 类。

code 6

private[netty] class Dispatcher(nettyEnv: NettyRpcEnv, numUsableCores: Int) extends Logging {
  ...
  private class EndpointData(
      val name: String,
      val endpoint: RpcEndpoint,
      val ref: NettyRpcEndpointRef) {
    val inbox = new Inbox(ref, endpoint)
  }
  // 记录了 name 与 EndpointData 的映射
  private val endpoints: ConcurrentMap[String, EndpointData] =
    new ConcurrentHashMap[String, EndpointData]
  
  // EndpointData 的 inbox 一旦接收到消息，就放入该队列
  private val receivers = new LinkedBlockingQueue[EndpointData]
  
  // post 消息
  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 {
        // 先将消息放入对应 EndpointData
        data.inbox.post(message)
        // 将 EndpointData 放入队列
        receivers.offer(data)
        None
      }
    }
    // We don't need to call `onStop` in the `synchronized` block
    error.foreach(callbackIfStopped)
  }
  
 /** Message loop used for dispatching messages. */
 private class MessageLoop extends Runnable {
    override def run(): Unit = {
      try {
        while (true) {
          try {
            val data = receivers.take()
            if (data == PoisonPill) {
              // Put PoisonPill back so that other MessageLoops can see it.
              receivers.offer(PoisonPill)
              return
            }
            // 调用 inbox 的处理方法。
            data.inbox.process(Dispatcher.this)
          } catch {
            case NonFatal(e) => logError(e.getMessage, e)
          }
        }
      } catch {
        case ie: InterruptedException => // exit
      }
    }
  }
  ...

所以我们能很容易的推想出 Inbox 的大致实现，首先他需要知道他的所负责的 Endpoint, 然后要有一个队列，记录发给他的消息，最后有一个批量处理队列中消息的方法(即 process)。在处理方法中，Inbox 根据消息的类型来执行Endpoint对应的操作，也就是启动，接受，停止，另外处理方法需要 Dispatcher 引用是为了在 EndPoint 被停止的时候从 Dispatcher 的列表中移出。代码就不贴了。

接下来的问题就是消息是如何交给 Dispatcher ，我们知道 RPC 服务端底层使用 TransportServer 来实现服务端，我们需要实现 RpcHandler 来定义消息处理。消息交给 Dispatcher 就是由 RpcHandler 的实现类 NettyRpcHandler 实现的，见 code 7：

code 7

private[netty] class NettyRpcHandler(
    dispatcher: Dispatcher,
    nettyEnv: NettyRpcEnv,
    streamManager: StreamManager) extends RpcHandler with Logging {
  ...
  override def receive(
      client: TransportClient,
      message: ByteBuffer,
      callback: RpcResponseCallback): Unit = {
    val messageToDispatch = internalReceive(client, message)
    dispatcher.postRemoteMessage(messageToDispatch, callback)
  }
  ...
}

很基本的操作，一个队列，多个线程去消费。就是好，就是块，厉害。

请求的响应

最最后就是请求的响应，分为两种情况，

• 一种如果是本地进程发送的请求，则响应的定义在 Promise 中，处理结束则调用 Promise 的对应方法；
• 另一种就是其他（远程）进程发送的请求，这是则需要利用网络传输模块来完成响应，即调用 code 7 中 receive 的 RpcResponseCallback 回调中对应的成功或失败的方法。

实战

待写，推荐直接看测试用例

总结

由上可见 RPC 模块比较重要的设计思路就是非阻塞批量的消息发送与接收。请求的响应与处理使用 Endpoint 来定义。在概念上更加明确，使用方式上更加灵活。

我们可以简单的画图总结 RPC 交互的整体过程：

待画