Spark RPC实现原理分析

Spark RPC模块架构图

Spark RPC是按照MailBox的设计思路来实现的，为了能够更直观地表达RPC的设计，我们先从RPC架构图来看，如下图所示：

该图主要描述了从客户端向服务器端发送远程消息和在服务器端发送本地消息的过程，该过程中显示了消息的流动及格式转换，还有各个组件之间的分工协作。

组件介绍

核心组件

Spark RPC通信主要有RpcEnv、RpcEndpoint、RpcEndpointRef这三个核心类。

1. RpcEndpoint

   该类定义了RPC通信过程中的服务器端对象，除了具有管理一组RpcEndpoint生命周期的操作（constructor -> onStart -> receive* -> onStop），并给出了通信过程中RpcEndpoint所具有的基于事件驱动的行为（连接、断开、网络异常），实际上对于Spark框架来说主要是接收消息并处理。

   private[spark] trait RpcEndpoint {
   
      /**
        * 当前RpcEndpoint所注册的[[RpcEnv]]
        */
      val rpcEnv: RpcEnv
   
      /**
        * 当前[[RpcEndpoint]]的代理，当`onStart`方法被调用时`self`生效，当`onStop`被调用时，`self`变成null。
        * 注意：在`onStart`方法被调用之前，[[RpcEndpoint]]对象还未进行注册，所以就没有有效的[[RpcEndpointRef]]。
        */
      final def self: RpcEndpointRef = {
          require(rpcEnv != null, "rpcEnv has not been initialized")
          rpcEnv.endpointRef(this)
      }
   
      /**
        * 用于处理从`RpcEndpointRef.send` 或 `RpcCallContext.reply`接收到的消息。
        * 如果接收到一个不匹配的消息，将会抛出SparkException异常，并发送给`onError`
        *
        * 通过上面的receive方法，接收由RpcEndpointRef.send方法发送的消息，
        * 该类消息不需要进行响应消息（Reply），而只是在RpcEndpoint端进行处理。
        */
      def receive: PartialFunction[Any, Unit] = {
          case _ => throw new SparkException(self + " does not implement 'receive'")
      }
   
      /**
        * 处理来自`RpcEndpointRef.ask`的消息，RpcEndpoint端处理完消息后，需要给调用RpcEndpointRef.ask的通信端返回响应消息。
        */
      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
      }
   
      /**
        * 当有远端连接到当前服务器时会被调用
        */
      def onConnected(remoteAddress: RpcAddress): Unit = {
          // By default, do nothing.
      }
   
      /**
        * 当远端与当前服务器断开时，该方法会被调用
        */
      def onDisconnected(remoteAddress: RpcAddress): Unit = {
          // By default, do nothing.
      }
   
      /**
        * 当前节点与远端之间的连接发生错误时，该方法将会被调用
        */
      def onNetworkError(cause: Throwable, remoteAddress: RpcAddress): Unit = {
          // By default, do nothing.
      }
   
      /**
        * 在 [[RpcEndpoint]] 开始处理消息之前被调用
        */
      def onStart(): Unit = {
          // By default, do nothing.
      }
   
      /**
        * 当[[RpcEndpoint]]正在停止时，该方法将会被调用。
        * `self`将会在该方法中被置位null，因此你不能使用它来发送消息。
        */
      def onStop(): Unit = {
          // By default, do nothing.
      }
   
      /**
        * A convenient method to stop [[RpcEndpoint]].
        */
      final def stop(): Unit = {
          val _self = self
          if (_self != null) {
              rpcEnv.stop(_self)
          }
      }
   }

2. RpcEndpointRef

   RpcEndpointRef是一个对RpcEndpoint的远程引用对象，通过它可以向远程的RpcEndpoint端发送消息以进行通信。RpcEndpointRef特质的定义，代码如下所示：

   private[spark] abstract class RpcEndpointRef(conf: SparkConf) extends Serializable with Logging {
   
      private[this] val maxRetries = RpcUtils.numRetries(conf)
      private[this] val retryWaitMs = RpcUtils.retryWaitMs(conf)
      private[this] val defaultAskTimeout = RpcUtils.askRpcTimeout(conf)
   
      /**
        * 返回[RpcEndpointRef]]的引用的远端服务器地址
        */
      def address: RpcAddress
   
      def name: String
   
      /**
        * 发送一条单向的异步消息，并且发送消息后不等待响应，亦即Send-and-forget。
        */
      def send(message: Any): Unit
   
      /**
        * 发送消息给相关的[[RpcEndpoint.receiveAndReply]]，并且返回一个 Future，能够在timeout时间内接收回复。
        * 该方法只会发送一次消息，失败后不重试。
        * 而ask方法发送消息后需要等待通信对端给予响应，通过Future来异步获取响应结果。
        */
      def ask[T: ClassTag](message: Any, timeout: RpcTimeout): Future[T]
   
      /**
        * 发送消息给相关的[[RpcEndpoint.receiveAndReply]]，并且返回一个 Future，能够在defaultAskTimeout时间内接收回复。
        * 该方法只会发送一次消息，失败后不重试。
        * 而ask方法发送消息后需要等待通信对端给予响应，通过Future来异步获取响应结果。
        */
      def ask[T: ClassTag](message: Any): Future[T] = ask(message, defaultAskTimeout)
   
      /**
        * 发送消息给相关的[[RpcEndpoint.receiveAndReply)]]，并且返回一个Future，能够在defaultAskTimeout时间内接收回复，如果超时则抛出异常。
        * 注意：该方法会阻塞当前线程，
        *
        * @param message the message to send
        * @tparam T type of the reply message
        * @return the reply message from the corresponding [[RpcEndpoint]]
        */
      def askSync[T: ClassTag](message: Any): T = askSync(message, defaultAskTimeout)
   
      /**
        * 发送消息给相关的[[RpcEndpoint.receiveAndReply)]]，并且返回一个Future，能够在timeout时间内接收回复，如果超时则抛出异常。
        * 注意：该方法会阻塞当前线程，
        *
        * @param 发送的消息内容
        * @param 超时时长
        * @tparam 响应消息的类型
        * @return 从[[RpcEndpoint]]端响应的消息内容
        */
      def askSync[T: ClassTag](message: Any, timeout: RpcTimeout): T = {
          val future = ask[T](message, timeout)
          timeout.awaitResult(future)
      }
   }
   

   上面的send方法的实现在NettyRpcEndpointRef中，表示发送消息后不等待响应，亦即Send-and-forget，具体实现如下：

   override def send(message: Any): Unit = {
    require(message != null, "Message is null")
    nettyEnv.send(new RequestMessage(nettyEnv.address /*如果是远程消息，则为null*/ , endpointRef, message))
   }

   可见，它是通过NettyRpcEnv来发送RequestMessage消息，并将当前NettyRpcEndpointRef封装到RequestMessage消息对象中发送出去，通信对端通过该NettyRpcEndpointRef能够识别出消息来源。

   而ask方法发送消息后需要等待通信对端给予响应，通过Future来异步获取响应结果，也是在NettyRpcEndpointRef中实现，如下所示：

   override def ask[T: ClassTag](message: Any, timeout: RpcTimeout): Future[T] = {
    nettyEnv.ask(new RequestMessage(nettyEnv.address /*如果是远程消息，则为null*/ , endpointRef, message), timeout)
   }

   类似的，也是通过NettyRpcEnv来发送一个RequestMessage消息。

3. RpcEnv

   一个RpcEnv是一个RPC环境对象，它负责管理RpcEndpoints的注册，以及如何从一个RpcEndpoint获取到一个RpcEndpointRef。RpcEndpoint是一个通信端，例如Spark集群中的Master，或Worker，都是一个RpcEndpoint。但是，如果想要与一个RpcEndpoint端进行通信，一定需要获取到该RpcEndpoint一个RpcEndpointRef，而获取该RpcEndpointRef只能通过一个RpcEnv环境对象来获取。

   private[spark] abstract class RpcEnv(conf: SparkConf) {
   
      private[spark] val defaultLookupTimeout = RpcUtils.lookupRpcTimeout(conf)
   
      /**
        * 返回已经注册的[[RpcEndpoint]]的RpcEndpointRef。
        * 该方法只用于[[RpcEndpoint.self]]方法实现中。
        * 如果终端相关的[[RpcEndpointRef]]不存在，则返回null。
        */
      private[rpc] def endpointRef(endpoint: RpcEndpoint): RpcEndpointRef
   
      /**
        * 如果是服务器模式，则返回当前服务器监听的地址；否则为空
        */
      def address: RpcAddress
   
      /**
        * 使用一个name来注册一个[[RpcEndpoint]]，并且返回它的[[RpcEndpointRef]]对象。
        * [[RpcEnv]]并不保证线程安全性。
        */
      def setupEndpoint(name: String, endpoint: RpcEndpoint): RpcEndpointRef
   
      /**
        * 通过一个URI来异步检索[[RpcEndpointRef]]对象
        */
      def asyncSetupEndpointRefByURI(uri: String): Future[RpcEndpointRef]
   
      /**
        * 通过一个URI来同步检索[[RpcEndpointRef]]对象
        */
      def setupEndpointRefByURI(uri: String): RpcEndpointRef = {
          defaultLookupTimeout.awaitResult(asyncSetupEndpointRefByURI(uri))
      }
   
      /**
        * 根据`address` 和 `endpointName`对 [[RpcEndpointRef]]进行同步检索。
        */
      def setupEndpointRef(address: RpcAddress, endpointName: String): RpcEndpointRef = {
          setupEndpointRefByURI(RpcEndpointAddress(address, endpointName).toString) // URI：
      }
   
      /**
        * 停止指定的[[RpcEndpoint]]对象。
        */
      def stop(endpoint: RpcEndpointRef): Unit
   
      /**
        * 异步关闭当前的[[RpcEnv]]。
        * 如果需要确保成功地退出[[RpcEnv]]，在执行[[shutdown()]]之后需要调用[[awaitTermination()]]。
        */
      def shutdown(): Unit
   
      /**
        * 等待直到[[RpcEnv]]退出。
        * TODO do we need a timeout parameter?
        */
      def awaitTermination(): Unit
   
      /**
        * 如果没有[[RpcEnv]]对象，那么[[RpcEndpointRef]]将不能被反序列化。
        * 因此，如果任何反序列化的对象中包含了[[RpcEndpointRef]]，那么这些反序列化的代码都应该在该方法中执行。
        */
      def deserialize[T](deserializationAction: () => T): T
   
      /**
        * 用于返回文件服务器的实例。
        * 如果RpcEnv不是以服务器模式运行，那么该项可能为null。
        *
        */
      def fileServer: RpcEnvFileServer
   
      /**
        * 打开一个通道从给定的URI下载文件。
        * 如果由RpcEnvFileServer返回的URI使用"spark"模式，那么该方法将会被工具类调用来进行文件检索。
        *
        * @param uri URI with location of the file.
        */
      def openChannel(uri: String): ReadableByteChannel
   
   }

收件箱InBox

从上面的体系结构图可知，InBox作用于服务器端。它与RpcEndpoint是一对一的关系，每一个命名唯一的RpcEndpoint对应一个线程安全的InBox。所有发送给一个RpcEndpoint的消息，都由对应的InBox进行存储。InBox提供一个process方法实现，该方法会在一个dispatcher-event-loop线程池中被调用，将InBox中的消息提供给关联的RpcEndpoint进行消费。

需要注意的是，如果通信端端点的实现是继承自ThreadSafeRpcEndpoint，则表明该Endpoint不允许并发处理消息。如果继承自RpcEndpoint，那么就可以并发的调用该服务。在具体的process方法中，如果enableConcurrent为false，即只允许单线程处理。那么执行process方法时，如果numActiveThreads大于0，说明已经至少有一个线程正在处理，则立即返回，取消本次处理操作。

Inbox.stop方法用于停止RpcEndpoint和InBox，具体实现如下：

def stop(): Unit = inbox.synchronized {
  // 该方法必须加锁，这样可以确保 "OnStop"是最后一条消息。
  if (!stopped) {
    enableConcurrent = false
    stopped = true
    messages.add(OnStop)
  }
}

它的原理是，首先该方法必须加锁inbox.synchronized，保证”OnStop”是消息列表中的最后一条消息。其次，为了保证能够安全地关闭RpcEndpoint，必须将enableConcurrent置位false，这样只有一个线程能够处理OnStop消息，并安全地执行RpcEndpoint.onStop方法并释放资源。当stop被执行后，其后向该InBox相关联的RpcEndpoint发送的消息将全部drop，并输出相应的日志记录。

消息转发路由Dispatcher

1. 提供RpcEndpoint注册

   在注册RpcEndpoint时，每个RpcEndpoint都需要有一个唯一的名称。在RpcEnv.setupEndpoint(name: String, endpoint: RpcEndpoint)方法的实现中就是直接调用registerRpcEndpoint进行端点注册的，并返回一个NettyRpcEndpointRef。

   def registerRpcEndpoint(name: String, endpoint: RpcEndpoint): NettyRpcEndpointRef = {
    val addr = RpcEndpointAddress(nettyEnv.address, name)
    val endpointRef = new NettyRpcEndpointRef(nettyEnv.conf, addr, nettyEnv)
    synchronized {
      if (stopped) {
        throw new IllegalStateException("RpcEnv has been stopped")
      }
   
      // 不能重复注册，根据RpcEndpoint的名称来检查唯一性
      if (endpoints.putIfAbsent(name, new EndpointData(name, endpoint, endpointRef)) != null) {
        throw new IllegalArgumentException(s"There is already an RpcEndpoint called $name")
      }
   
      val data = endpoints.get(name)
      endpointRefs.put(data.endpoint, data.ref)
      // 向receivers添加当前的Endpoint，receivers表示当前有未读消息的用户列表，
      receivers.offer(data) // for the OnStart message
    }
    endpointRef
   }

   • EndpointData是一个简单的JavaBean类，用于封装RpcEndpoint相关的组件和属性：

     /**
     * 一个JavaBean对象，用于封装Endpoint对象、Endpoint对象唯一名称、EndpointRef以及收件箱
     * 
     * @param name     Endpoint name
     * @param endpoint 服务器端
     * @param ref      服务器端引用
     */
     private class EndpointData(val name: String, val endpoint: RpcEndpoint, val ref: NettyRpcEndpointRef) {
     val inbox = new Inbox(ref, endpoint)
     }

   • endpoints

     该属性用于维护RpcEndpoint的名称与EndpointData之间的映射关系\

发件箱OutBox

从上面的体系结构图可知，OnBox作用于客户端。类似于收件箱InBox，它与TransportClient是一对一的关系，而一个TransportClient对应着一个远程的RpcEndPoint。OutBox提供的主要功能如下：

1. 使用send方法将消息投递到发件箱

   OutBox提供send方法，会将RpcEnv中所有转发给某个RpcEndPoint的消息都先放到一个messages链表中，最后就开始使用drainOutbox来清空发件箱。

   /**
    * 用于发送消息。
    * - 如果目前没有可用的连接，则将消息缓存并建立一个连接。
    * - 如果[[Outbox]]已经停止，那么sender将会抛出一个[[SparkException]]
    */
   def send(message: OutboxMessage): Unit = {
    val dropped = synchronized {
      if (stopped) {
        true
      } else {
        messages.add(message)
        false
      }
    }
    if (dropped) {
      // 如果[[Outbox]]已经停止，那么sender将会抛出一个[[SparkException]]
      message.onFailure(new SparkException("Message is dropped because Outbox is stopped"))
    } else {
      drainOutbox()
    }
   }

2. 使用drainOutbox清空发件箱

   drainOutbox主要用于清空发件箱中的消息，消息会通过传输层TransportClient发送给远端服务器。该方法会在开始处进行一系列的检查，需要保证传输层的连接已经建立，如果没有建立，则向nettyEnv.clientConnectionExecutor提交建立连接的任务后并返回，待连接任务完成后会再次调用drainOutox方法。另外，drainOutox会保证线程安全性，通过布尔值draining可以保证同一时刻只会有一个线程能够进行消息的处理和发送。

   具体的消息发送逻辑则交由OutboxMessage的实现来完成，OutboxMessage有两个子类，OneWayOutboxMessage和RpcOutboxMessage，分别对应调用RpcEndpoint的receive和receiveAndReply方法。

   /**
    * sealed作用：
    * 1. 其修饰的trait，class只能在当前文件里面被继承
    * 2. 用sealed修饰这样做的目的是告诉scala编译器在检查模式匹配的时候，让scala知道这些case的所有情况，
    *    scala就能够在编译的时候进行检查，看你写的代码是否有没有漏掉什么没case到，减少编程的错误。
    */
   private[netty] sealed trait OutboxMessage {
   
      def sendWith(client: TransportClient): Unit
   
      def onFailure(e: Throwable): Unit
   
   }

   对于OneWayOutboxMessage，由于不需要返回值，则简单地通过调用传输层client.send方法将消息发出。

   private[netty] case class OneWayOutboxMessage(content: ByteBuffer) extends OutboxMessage with Logging {
   
      override def sendWith(client: TransportClient): Unit = {
          client.send(content)
      }
   
      override def onFailure(e: Throwable): Unit = {
          e match {
              case e1: RpcEnvStoppedException => logWarning(e1.getMessage)
              case e1: Throwable => logWarning(s"Failed to send one-way RPC.", e1)
          }
      }
   
   }

   对于RpcOutboxMessage，由于需要服务器的响应结果，因此需要实现传输层提供的RpcResponseCallback接口，并提供onFailure和onSuccess的方法实现。在实际的发送消息时会使用client.sendRpc方法，将消息内容和RpcResponseCallback对象传递给传输层，该方法会立即返回一个requestId。

   而传输层底层会有独立的线程负责将消息序列化并且发送出去，每个Message都会返回一个UUID，由底层来维护一个发送出去消息与其Callback的HashMap。

   • 如果请求超时，会通过requestId在传输层中移除该RPC请求，从而达到取消消息发送的效果；
   • 如果请求的消息成功返回，则会使用RpcResponseCallback对象根据返回的状态回调对应的onFailure和onSuccess的方法，进而回调Spark core中的业务逻辑，执行Promise/Future的done方法，上层退出阻塞。

   private[netty] case class RpcOutboxMessage(content: ByteBuffer,
                                             _onFailure: (Throwable) => Unit,
                                             _onSuccess: (TransportClient, ByteBuffer) => Unit)
          extends OutboxMessage with RpcResponseCallback with Logging {
   
      private var client: TransportClient = _
      private var requestId: Long = _
   
      override def sendWith(client: TransportClient): Unit = {
          this.client = client
          this.requestId = client.sendRpc(content, this)
      }
   
      def onTimeout(): Unit = {
          if (client != null) {
              client.removeRpcRequest(requestId)
          } else {
              logError("Ask timeout before connecting successfully")
          }
      }
   
      override def onFailure(e: Throwable): Unit = {
          _onFailure(e)
      }
   
      override def onSuccess(response: ByteBuffer): Unit = {
          _onSuccess(client, response)
      }
   
   }

3. 关闭收件箱

   网络连接错误和RpcEnv的停止运行都会触发OutBox的关闭和资源的清理，OutBox关闭的处理逻辑如下：

   • 如果connectFuture不为空，说明这会正在执行连接任务，那么调用connectFuture.cancel(true)方法，将任务取消。

   • 调用closeClient方法，关闭客户端，这里仅仅将client引用置为null，但并不是真正的关闭，因为需要重用连接。

   • 调用nettyEnv.removeOutbox(remoteAddress)方法，从nettyEnv中移除OutBox，因此将来的消息将会使用一个新的或原有的client连接并创建一个新的OutBox。
   • 执行所有还未处理的消息的onFailure方法，并告知失败的原因。

Spark RPC服务的流程

使用示例

package org.apache.spark.rpc

import org.apache.spark.rpc.netty.NettyRpcEnvFactory
import org.apache.spark.{SecurityManager, SparkConf, SparkFunSuite}
import org.scalatest.BeforeAndAfterAll
import org.scalatest.concurrent.Eventually.{eventually, interval, timeout}
import scala.concurrent.duration._
import scala.language.postfixOps

class SparkRPCTest extends SparkFunSuite with BeforeAndAfterAll {

    def createRpcEnv(conf: SparkConf, name: String, port: Int, clientMode: Boolean = false): RpcEnv = {
        val config = RpcEnvConfig(conf, name, "localhost", "localhost", port, new SecurityManager(conf), clientMode)
        new NettyRpcEnvFactory().create(config)
    }

    test("send a message to server from local") {
        var serverRpcEnv: RpcEnv = null
        try {
            @volatile var serverReceivedMsg: String = null
            serverRpcEnv = createRpcEnv(new SparkConf(), "server", 0)
            val serverEndpointRef = serverRpcEnv.setupEndpoint("server-endpoint", new RpcEndpoint {
                override val rpcEnv: RpcEnv = serverRpcEnv

                override def receive: PartialFunction[Any, Unit] = {
                    case msg: String => serverReceivedMsg = msg
                }
            })
            serverEndpointRef.send("hello")
            eventually(timeout(5 seconds), interval(10 millis)) {
                assert("hello" === serverReceivedMsg)
            }
        } finally {
            destory(serverRpcEnv)
        }
    }

    test("send a message to server from client") {
        var serverRpcEnv: RpcEnv = null
        var clientRpcEnv: RpcEnv = null
        try {
            @volatile var serverReceivedMsg: String = null
            serverRpcEnv = createRpcEnv(new SparkConf(), "server", 0)
            serverRpcEnv.setupEndpoint("server-endpoint", new RpcEndpoint {
                override val rpcEnv: RpcEnv = serverRpcEnv

                override def receive: PartialFunction[Any, Unit] = {
                    case msg: String => serverReceivedMsg = msg
                }
            })
            clientRpcEnv = createRpcEnv(new SparkConf(), "client", 0, clientMode = true)
            val serverEndpointRef = clientRpcEnv.setupEndpointRef(serverRpcEnv.address, "server-endpoint")
            serverEndpointRef.send("hello")
            eventually(timeout(5 seconds), interval(10 millis)) {
                assert("hello" === serverReceivedMsg)
            }
        } finally {
            destory(clientRpcEnv)
            destory(serverRpcEnv)
        }
    }

    def destory(env: RpcEnv): Unit = {
        try {
            if (env != null) {
                env.shutdown()
                env.awaitTermination()
            }
        } finally {
            super.afterAll()
        }
    }
}

服务器启动

服务器响应

服务器响应主要分为两个阶段，

1. 第一阶段，IO接收

   TransportRequestHandler是netty的回调handler，它会将请求最终转发给内部NettyRpcHandler的receive方法进行处理，方法内部会调用Dispatcher将RpcMessage投递到对应的Inbox中，到此结束。

2. 第二阶段，IO响应

   MessageLoop会获取待处理的RpcMessage进行处理，实际上就是调用RpcEndpoint的业务逻辑得到对应的结果，之后通过回调NettyRpcCallContext.reply方法，将响应结果序列化后，再通过回调RpcResponseCallback.onSuccess方法将序列化后的结果传给TransportRequestHandler.this.respond方法，该方法会将执行结果封装为RpcResponse后通过Channel发送回客户端。

客户端请求

1. 第一阶段，IO发送

   利用RpcEndpointRef做send或者ask动作。这里以send为例，send会先进行消息的序列化，然后投递到指定地址的Outbox中，Outbox如果发现连接未建立则先尝试建立连接，然后调用底层的TransportClient发送数据，直接通过该netty的API完成，完成后即可返回。这里返回了UUID作为消息的标识，用于下一个阶段的回调，使用的角度来说可以返回一个Future，客户端可以阻塞或者继续做其他操作。

2. 第二阶段，IO接收

   TransportResponseHandler接收到远程的响应后，会根据RpcResponse中的requestId获得RpcResponseCallback，然后根据回调对应的onSuccess方法。接着将执行结果反序列化后，回调第一阶段的Future，完成调用。