Spark RPC 学习笔记

2020-12-23

前些天发现了一个巨牛的人工智能学习网站，通俗易懂，风趣幽默，忍不住给大家分享一下。点击跳转到网站：https://www.captainai.net/dongkelun

欢迎关注我的公众号

前言

学习记录一下Spark源码中的RPC，本文基于Spark3源码进行学习研究

三个重要的类

RpcEnv、RpcEndpoint、RpcEndpointRef主要是想搞懂这三个之间的关系

子类

Spark 1.6.3中默认使用了Netty作为底层的实现，但Akka的依赖依然存在；而Spark 2.1.0中的底层实现只有Netty，这样用户可以方便的使用不同版本的Akka或者将来某种更好的底层实现

NettyRpcEnv

private[netty] class NettyRpcEnv(
    val conf: SparkConf,
    javaSerializerInstance: JavaSerializerInstance,
    host: String,
    securityManager: SecurityManager,
    numUsableCores: Int) extends RpcEnv(conf) with Logging {

NettyRpcEndpointRef

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

RpcEndpoint 有两个子类trait（特质）


private[spark] trait ThreadSafeRpcEndpoint extends RpcEndpoint


private[spark] trait IsolatedRpcEndpoint extends RpcEndpoint {

  def threadCount(): Int = 1

}

真正的实现是在具体的内部类里，比如
CoarseGrainedSchedulerBackend 里面的 DriverEndpoint

class DriverEndpoint extends IsolatedRpcEndpoint with Logging {

    override val rpcEnv: RpcEnv = CoarseGrainedSchedulerBackend.this.rpcEnv

    protected val addressToExecutorId = new HashMap[RpcAddress, String]

    // Spark configuration sent to executors. This is a lazy val so that subclasses of the
    // scheduler can modify the SparkConf object before this view is created.
    private lazy val sparkProperties = scheduler.sc.conf.getAll
      .filter { case (k, _) => k.startsWith("spark.") }
      .toSeq

    private val logUrlHandler: ExecutorLogUrlHandler = new ExecutorLogUrlHandler(
      conf.get(UI.CUSTOM_EXECUTOR_LOG_URL))

    override def onStart(): Unit = {
        ......

也可以是匿名内部类如在测试RpcEnvSuite里的

val rpcEndpointRef = env.setupEndpoint("send-locally", new RpcEndpoint {
      override val rpcEnv = env

      override def receive = {
        case msg: String => {
          println(msg)
          message = msg
        }
      }
    })

注册流程

创建RpcEndpoint(子类)

1 2	new DriverEndpoint() new RpcEndpoint{...

创建RpcEnv

实际上是通过NettyRpcEnvFactory的create方法创建并返回NettyRpcEnv

SparkContext初始化时：
调用顺序createSparkEnv->SparkEnv.createDriverEnv->create->RpcEnv.create

 // Create the Spark execution environment (cache, map output tracker, etc)
_env = createSparkEnv(_conf, isLocal, listenerBus)

  private[spark] def createSparkEnv(
      conf: SparkConf,
      isLocal: Boolean,
      listenerBus: LiveListenerBus): SparkEnv = {
    SparkEnv.createDriverEnv(conf, isLocal, listenerBus, SparkContext.numDriverCores(master, conf))
  }

  /**
   * Create a SparkEnv for the driver.
   */
  private[spark] def createDriverEnv(
      conf: SparkConf,
      isLocal: Boolean,
      listenerBus: LiveListenerBus,
      numCores: Int,
      mockOutputCommitCoordinator: Option[OutputCommitCoordinator] = None): SparkEnv = {
    assert(conf.contains(DRIVER_HOST_ADDRESS),
      s"${DRIVER_HOST_ADDRESS.key} is not set on the driver!")
    assert(conf.contains(DRIVER_PORT), s"${DRIVER_PORT.key} is not set on the driver!")
    val bindAddress = conf.get(DRIVER_BIND_ADDRESS)
    val advertiseAddress = conf.get(DRIVER_HOST_ADDRESS)
    val port = conf.get(DRIVER_PORT)
    val ioEncryptionKey = if (conf.get(IO_ENCRYPTION_ENABLED)) {
      Some(CryptoStreamUtils.createKey(conf))
    } else {
      None
    }
    create(
      conf,
      SparkContext.DRIVER_IDENTIFIER,
      bindAddress,
      advertiseAddress,
      Option(port),
      isLocal,
      numCores,
      ioEncryptionKey,
      listenerBus = listenerBus,
      mockOutputCommitCoordinator = mockOutputCommitCoordinator
    )
  }
  private def create(
      conf: SparkConf,
      executorId: String,
      bindAddress: String,
      advertiseAddress: String,
      port: Option[Int],
      isLocal: Boolean,
      numUsableCores: Int,
      ioEncryptionKey: Option[Array[Byte]],
      listenerBus: LiveListenerBus = null,
      mockOutputCommitCoordinator: Option[OutputCommitCoordinator] = None): SparkEnv = {
      ...
    val rpcEnv = RpcEnv.create(systemName, bindAddress, advertiseAddress, port.getOrElse(-1), conf,
      securityManager, numUsableCores, !isDriver)
      ...
    val envInstance = new SparkEnv(
      executorId,
      rpcEnv,
      serializer,
      closureSerializer,
      serializerManager,
      mapOutputTracker,
      shuffleManager,
      broadcastManager,
      blockManager,
      securityManager,
      metricsSystem,
      memoryManager,
      outputCommitCoordinator,
      conf)
      ...
      envInstance

RpcEnv.create方法：

def create(
      name: String,
      bindAddress: String,
      advertiseAddress: String,
      port: Int,
      conf: SparkConf,
      securityManager: SecurityManager,
      numUsableCores: Int,
      clientMode: Boolean): RpcEnv = {
    val config = RpcEnvConfig(conf, name, bindAddress, advertiseAddress, port, securityManager,
      numUsableCores, clientMode)
    new NettyRpcEnvFactory().create(config)
  }

实际调用的NettyRpcEnvFactory的create方法返回NettyRpcEnv

  def create(config: RpcEnvConfig): RpcEnv = {
    val sparkConf = config.conf
    // Use JavaSerializerInstance in multiple threads is safe. However, if we plan to support
    // KryoSerializer in future, we have to use ThreadLocal to store SerializerInstance
    val javaSerializerInstance =
      new JavaSerializer(sparkConf).newInstance().asInstanceOf[JavaSerializerInstance]
    val nettyEnv =
      new NettyRpcEnv(sparkConf, javaSerializerInstance, config.advertiseAddress,
        config.securityManager, config.numUsableCores)
    if (!config.clientMode) {
      val startNettyRpcEnv: Int => (NettyRpcEnv, Int) = { actualPort =>
        nettyEnv.startServer(config.bindAddress, actualPort)
        (nettyEnv, nettyEnv.address.port)
      }
      try {
        Utils.startServiceOnPort(config.port, startNettyRpcEnv, sparkConf, config.name)._1
      } catch {
        case NonFatal(e) =>
          nettyEnv.shutdown()
          throw e
      }
    }
    nettyEnv
  }
}

下面是测试类RpcEnvSuite的创建方法

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

向RpcEnv注册RpcEndpoint并返回RpcEndpointRef

1	rpcEnv.setupEndpoint(ENDPOINT_NAME, createDriverEndpoint())

实际上是调用NettyRpcEnv的setupEndpoint方法

1
2
3

override def setupEndpoint(name: String, endpoint: RpcEndpoint): RpcEndpointRef = {
   dispatcher.registerRpcEndpoint(name, endpoint)
 }

然后调用dispatcher.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")
    }
    if (endpoints.containsKey(name)) {
      throw new IllegalArgumentException(s"There is already an RpcEndpoint called $name")
    }

    // This must be done before assigning RpcEndpoint to MessageLoop, as MessageLoop sets Inbox be
    // active when registering, and endpointRef must be put into endpointRefs before onStart is
    // called.
    endpointRefs.put(endpoint, endpointRef)

    var messageLoop: MessageLoop = null
    try {
      messageLoop = endpoint match {
        case e: IsolatedRpcEndpoint =>
          new DedicatedMessageLoop(name, e, this)
        case _ =>
          sharedLoop.register(name, endpoint)
          sharedLoop
      }
      endpoints.put(name, messageLoop)
    } catch {
      case NonFatal(e) =>
        endpointRefs.remove(endpoint)
        throw e
    }
  }
  endpointRef
}

RpcEndpointRef.send

实际上调用NettyRpcEndpointRef.send

override def send(message: Any): Unit = {
  require(message != null, "Message is null")
  nettyEnv.send(new RequestMessage(nettyEnv.address, this, message))
}

然后调用NettyRpcEnv.send

private[netty] def send(message: RequestMessage): Unit = {
    val remoteAddr = message.receiver.address
    if (remoteAddr == address) {
      // Message to a local RPC endpoint.
      try {
        dispatcher.postOneWayMessage(message)
      } catch {
        case e: RpcEnvStoppedException => logDebug(e.getMessage)
      }
    } else {
      // Message to a remote RPC endpoint.
      postToOutbox(message.receiver, OneWayOutboxMessage(message.serialize(this)))
    }
  }

这里会根据address判断是否本地还是远程调用，下面按本地模式进行研究，即调用Dispatcher的postOneWayMessage方法


/** Posts a one-way message. */
def postOneWayMessage(message: RequestMessage): Unit = {
  postMessage(message.receiver.name, OneWayMessage(message.senderAddress, message.content),
    (e) => e match {
      // SPARK-31922: in local cluster mode, there's always a RpcEnvStoppedException when
      // stop is called due to some asynchronous message handling. We catch the exception
      // and log it at debug level to avoid verbose error message when user stop a local
      // cluster in spark shell.
      case re: RpcEnvStoppedException => logDebug(s"Message $message dropped. ${re.getMessage}")
      case _ => throw e
    })
}

private def postMessage(
    endpointName: String,
    message: InboxMessage,
    callbackIfStopped: (Exception) => Unit): Unit = {
  val error = synchronized {
    val loop = endpoints.get(endpointName)
    if (stopped) {
      Some(new RpcEnvStoppedException())
    } else if (loop == null) {
      Some(new SparkException(s"Could not find $endpointName."))
    } else {
      loop.post(endpointName, message)
      None
    }
  }
  // We don't need to call `onStop` in the `synchronized` block
  error.foreach(callbackIfStopped)
}

这里的loop是在上面的registerRpcEndpoint方法中根据判断endpoint是否是IsolatedRpcEndpoint的子类来区分成两种

var messageLoop: MessageLoop = null
try {
  messageLoop = endpoint match {
    case e: IsolatedRpcEndpoint =>
      new DedicatedMessageLoop(name, e, this)
    case _ =>
      sharedLoop.register(name, endpoint)
      sharedLoop
  }
  endpoints.put(name, messageLoop)

这里先按sharedLoop来进行分析，其中

1	private lazy val sharedLoop = new SharedMessageLoop(nettyEnv.conf, this, numUsableCores)

loop.post:

override def post(endpointName: String, message: InboxMessage): Unit = {
  val inbox = endpoints.get(endpointName)
  inbox.post(message)
  setActive(inbox)
}

向inbox插入一条InboxMessage，并将该inbox设置为active

receive

SharedMessageLoop中

override protected val threadpool: ThreadPoolExecutor = {
  val numThreads = getNumOfThreads(conf)
  val pool = ThreadUtils.newDaemonFixedThreadPool(numThreads, "dispatcher-event-loop")
  for (i <- 0 until numThreads) {
    pool.execute(receiveLoopRunnable)
  }
  pool
}

1
2
3

protected val receiveLoopRunnable = new Runnable() {
   override def run(): Unit = receiveLoop()
 }

  private def receiveLoop(): Unit = {
    try {
      while (true) {
        try {
          val inbox = active.take()
          if (inbox == MessageLoop.PoisonPill) {
            // Put PoisonPill back so that other threads can see it.
            setActive(MessageLoop.PoisonPill)
            return
          }
          inbox.process(dispatcher)
        } catch {
          case NonFatal(e) => logError(e.getMessage, e)
        }
      }
    } catch {
      case _: InterruptedException => // exit
        case t: Throwable =>
          try {
            // Re-submit a receive task so that message delivery will still work if
            // UncaughtExceptionHandler decides to not kill JVM.
            threadpool.execute(receiveLoopRunnable)
          } finally {
            throw t
          }
    }
  }
}

循环调用inbox.process


/**
 * Process stored messages.
 */
def process(dispatcher: Dispatcher): Unit = {
  var message: InboxMessage = null
  inbox.synchronized {
    if (!enableConcurrent && numActiveThreads != 0) {
      return
    }
    message = messages.poll()
    if (message != null) {
      numActiveThreads += 1
    } else {
      return
    }
  }
  while (true) {
    safelyCall(endpoint) {
      message match {
        case RpcMessage(_sender, content, context) =>
          try {
            endpoint.receiveAndReply(context).applyOrElse[Any, Unit](content, { msg =>
              throw new SparkException(s"Unsupported message $message from ${_sender}")
            })
          } catch {
            case e: Throwable =>
              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
          }

        case OneWayMessage(_sender, content) =>
          endpoint.receive.applyOrElse[Any, Unit](content, { msg =>
            throw new SparkException(s"Unsupported message $message from ${_sender}")
          })

        case OnStart =>
          endpoint.onStart()
          if (!endpoint.isInstanceOf[ThreadSafeRpcEndpoint]) {
            inbox.synchronized {
              if (!stopped) {
                enableConcurrent = true
              }
            }
          }

        case OnStop =>
          val activeThreads = inbox.synchronized { inbox.numActiveThreads }
          assert(activeThreads == 1,
            s"There should be only a single active thread but found $activeThreads threads.")
          dispatcher.removeRpcEndpointRef(endpoint)
          endpoint.onStop()
          assert(isEmpty, "OnStop should be the last message")

        case RemoteProcessConnected(remoteAddress) =>
          endpoint.onConnected(remoteAddress)

        case RemoteProcessDisconnected(remoteAddress) =>
          endpoint.onDisconnected(remoteAddress)

        case RemoteProcessConnectionError(cause, remoteAddress) =>
          endpoint.onNetworkError(cause, remoteAddress)
      }
    }

    inbox.synchronized {
      // "enableConcurrent" will be set to false after `onStop` is called, so we should check it
      // every time.
      if (!enableConcurrent && numActiveThreads != 1) {
        // If we are not the only one worker, exit
        numActiveThreads -= 1
        return
      }
      message = messages.poll()
      if (message == null) {
        numActiveThreads -= 1
        return
      }
    }
  }
}

inbox.process会匹配消息类型，然后根据消息类型，执行endpoint.receive方法，这里的消息类型是OneWayMessage即

case OneWayMessage(_sender, content) =>
endpoint.receive.applyOrElse[Any, Unit](content, { msg =>
  throw new SparkException(s"Unsupported message $message from ${_sender}")
})

最后看一下DedicatedMessageLoop，其实DedicatedMessageLoop和SharedMessageLoop差不多

private val inbox = new Inbox(name, endpoint)
// Mark active to handle the OnStart message.
setActive(inbox)
override def post(endpointName: String, message: InboxMessage): Unit = {
  require(endpointName == name)
  inbox.post(message)
  setActive(inbox)
}

同样有一个线程池去调用receiveLoopRunnable,下面的就和SharedMessageLoop一样了

override protected val threadpool = if (endpoint.threadCount() > 1) {
  ThreadUtils.newDaemonCachedThreadPool(s"dispatcher-$name", endpoint.threadCount())
} else {
  ThreadUtils.newDaemonSingleThreadExecutor(s"dispatcher-$name")
}

(1 to endpoint.threadCount()).foreach { _ =>
  threadpool.submit(receiveLoopRunnable)
}

onStart

这里提一下onStart方法，因为注释中写到

 * An end point for the RPC that defines what functions to trigger given a message.
 *
 * It is guaranteed that `onStart`, `receive` and `onStop` will be called in sequence.
 *
 * The life-cycle of an endpoint is:
 *
 * {@code constructor -> onStart -> receive* -> onStop}
......
   /**
   * Invoked before [[RpcEndpoint]] starts to handle any message.
   */
  def onStart(): Unit = {
    // By default, do nothing.
  }

即一个RpcEndpoint生命周期：构建->onStart→receive→onStop，那么onstart方法是在receive方法之前先执行的，但是在代码这个类中无法看出如何最新执行，且网上的
博客也没有写清楚如何执行，然后我就带着以为仔细调试和阅读了一下源码，发现是这样的：

首先在上面的inbox初始化时，首先会执行下面的代码

 @GuardedBy("this")
 protected val messages = new java.util.LinkedList[InboxMessage]()

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

那么在上面提到的process方法中messages会首先取出OnStart，而在模式匹配时匹配到OnStart，就会执行OnStart方法

message = messages.poll()

  case OnStart =>
    endpoint.onStart()
    if (!endpoint.isInstanceOf[ThreadSafeRpcEndpoint]) {
      inbox.synchronized {
        if (!stopped) {
          enableConcurrent = true
        }
      }
    }

在CoarseGrainedSchedulerBackend里的DriverEndpoint就重写了onStart方法来提前执行一些准备工作

override def onStart(): Unit = {
  // Periodically revive offers to allow delay scheduling to work
  // 调度程序为了运行任务而重新提供work资源的间隔长度。
  val reviveIntervalMs = conf.get(SCHEDULER_REVIVE_INTERVAL).getOrElse(1000L)
  //每隔1秒，给自己发一个ReviveOffers,发给receive函数
  reviveThread.scheduleAtFixedRate(() => Utils.tryLogNonFatalError {
    Option(self).foreach(_.send(ReviveOffers))
  }, 0, reviveIntervalMs, TimeUnit.MILLISECONDS)
}

测试Demo

import org.apache.spark.{SecurityManager, SparkConf, SparkEnv}
import org.apache.spark.rpc.netty.NettyRpcEnvFactory
import org.scalatest.concurrent.Eventually.{eventually, interval, timeout}

import scala.concurrent.duration._

/**
 * Created by dongkelun on 2020/12/24 11:20
 */
object RpcTestMain {
  def main(args: Array[String]): Unit = {
    val conf = new SparkConf()
    val env = createRpcEnv(conf, "local", 0)
    @volatile var message: String = null
    val rpcEndpointRef = env.setupEndpoint("send-locally", new RpcEndpoint {
      override val rpcEnv = env

      override def onStart(): Unit = {
        println("start hello endpoint")
      }

      override def receive = {
        case msg: String => {
          println(msg)
          message = msg
        }
      }
    })
    rpcEndpointRef.send("hello")

    eventually(timeout(5.seconds), interval(10.milliseconds)) {
      assert("hello" == message)
    }

    if (env != null) {
      env.shutdown()
    }
    SparkEnv.set(null)

  }

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

}

运行结果：

1 2	start hello endpoint hello

本文由 董可伦 发表于伦少的博客 ,采用署名-非商业性使用-禁止演绎 3.0进行许可。

非商业转载请注明作者及出处。商业转载请联系作者本人。

本文标题：Spark RPC 学习笔记

本文链接：https://dongkelun.com/2020/12/23/sparkRPC/

欢迎关注我的公众号