Spark源码阅读篇-Rpc通信-MessageLoop消息循环

上一节介绍了Dispatcher消息分发器的源码，里面有一个sharedLoop是共享消息循环，主要是负责保存多个RpcEndpoint端点的消息，这一节就来结合源码看一下消息循环是怎么保存消息并且发送给具体的端点的。

MessageLoop源码

//消息循环被Dispatcher用来给多个节点传递消息 sealed--限制继承层级
private sealed abstract class MessageLoop(dispatcher: Dispatcher) extends Logging {

  //基于单向链表实现的线程安全的阻塞队列 每个结点是一个inbox
  private val active = new LinkedBlockingQueue[Inbox]()

  //消息循环任务 主要是通过线程运行receiveLoop方法，检查是否收到新消息，根据收到的消息应对处理
  //该方法定时运行
  protected val receiveLoopRunnable = new Runnable() {
    override def run(): Unit = receiveLoop()
  }

  //线程池
  protected val threadpool: ExecutorService
  
  //表示MessageLoop状态 
  private var stopped = false
  
  //向指定的endpoint节点的inbox发消息
  def post(endpointName: String, message: InboxMessage): Unit

  //注销某个节点
  def unregister(name: String): Unit

  //将MessageLoop停止
  def stop(): Unit = {
    synchronized {
      if (!stopped) {
        //在消息循环中放入一个毒箱 当线程取到毒丸 则停止处理消息
        setActive(MessageLoop.PoisonPill)
        //线程池停止
        threadpool.shutdown()
        stopped = true
      }
    }
    //线程池等待
    threadpool.awaitTermination(Long.MaxValue, TimeUnit.MILLISECONDS)
  }

  //向active末尾放入一个inbox收件箱 offer表示在队列末尾放入一个结点
  protected final def setActive(inbox: Inbox): Unit = active.offer(inbox)

  //在线程中不间断执行，收到消息则放入到active队列末尾
  private def receiveLoop(): Unit = {
    try {
      while (true) {
        try {
          //取active队列第一个结点
          val inbox = active.take()
          if (inbox == MessageLoop.PoisonPill) {//如果取到毒丸 就放回队列末尾
            setActive(MessageLoop.PoisonPill)
            return
          }
          //inbox针对不同类型的消息，调用endpoint对应的方法来处理，inbox与endpoint是一一对应的
          inbox.process(dispatcher)
        } catch {
          case NonFatal(e) => logError(e.getMessage, e)
        }
      }
    } catch {
      case _: InterruptedException => // exit
        case t: Throwable =>
          try {
            //重新提交接收任务，以便消息传递仍然可以工作
            //UncaughtExceptionHandler决定不杀死JVM。
            threadpool.execute(receiveLoopRunnable)
          } finally {
            throw t
          }
    }
  }
}

private object MessageLoop {
  //毒丸，当线程取到毒丸的时候需要停止处理消息
  val PoisonPill = new Inbox(null, null)
}

//共享消息循环 同时给多个endpoint节点传递消息 归属于Dispatcher
private class SharedMessageLoop(
    conf: SparkConf,
    dispatcher: Dispatcher,
    numUsableCores: Int)
  extends MessageLoop(dispatcher) {
  
  //endpoints用来表示节点和inbox的映射关系 通常一个endpoint对应一个inbox
  private val endpoints = new ConcurrentHashMap[String, Inbox]()

  //获取线程数
  private def getNumOfThreads(conf: SparkConf): Int = {
    //可用核数 即当前Dispatcher分配了多少核数
    val availableCores =
      if (numUsableCores > 0) numUsableCores else Runtime.getRuntime.availableProcessors()
    //线程数需要大于或者等于2
    val modNumThreads = conf.get(RPC_NETTY_DISPATCHER_NUM_THREADS)
      .getOrElse(math.max(2, availableCores))
    //根绝所属的节点的角色获取配置信息 driver和executor默认线程数不一样
    conf.get(EXECUTOR_ID).map { id =>
      val role = if (id == SparkContext.DRIVER_IDENTIFIER) "driver" else "executor"
      conf.getInt(s"spark.$role.rpc.netty.dispatcher.numThreads", modNumThreads)
    }.getOrElse(modNumThreads)
  }

  //用来扫描消息循环并发送消息的线程池
  override protected val threadpool: ThreadPoolExecutor = {
    //获取默认线程数
    val numThreads = getNumOfThreads(conf)
    //创建线程池
    val pool = ThreadUtils.newDaemonFixedThreadPool(numThreads, "dispatcher-event-loop")
    //启动线程池中所有线程 都开始读取active队列处理消息
    for (i <- 0 until numThreads) {
      pool.execute(receiveLoopRunnable)
    }
    pool
  }

  //给指定的节点endpoint发消息
  override def post(endpointName: String, message: InboxMessage): Unit = {
    //获取该节点对应的inbox-收件箱
    val inbox = endpoints.get(endpointName)
    //inbox将消息放到队列末尾，同样是线程扫描处理
    inbox.post(message)
    //将inbox放到active队列末尾
    setActive(inbox)
  }

  //在map中删除节点endpoint和inbox的对应关系
  override def unregister(name: String): Unit = {
    val inbox = endpoints.remove(name)
    if (inbox != null) {
      //停止inbox 
      inbox.stop()
      // Mark active to handle the OnStop message.
      //将inbox状态设置为停止
      setActive(inbox)
    }
  }

  //将节点endpoint和inbox的对应关系放入map注册
  def register(name: String, endpoint: RpcEndpoint): Unit = {
    val inbox = new Inbox(name, endpoint)
    endpoints.put(name, inbox)
    //启动inbox
    setActive(inbox)
  }
}

//单个节点独享的消息循环
private class DedicatedMessageLoop(
    name: String,
    endpoint: IsolatedRpcEndpoint,//独立节点
    dispatcher: Dispatcher)
  extends MessageLoop(dispatcher) {

  //独立节点的收件箱inbox
  private val inbox = new Inbox(name, endpoint)

  //独立节点的线程池 处理消息
  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)
  }

  //启动inbox 实际上是启动对应的endpoint
  setActive(inbox)

  //发送消息给独立节点
  override def post(endpointName: String, message: InboxMessage): Unit = {
    //如果目标节点是独立节点则发送消息 否则抛出异常 require
    require(endpointName == name)
    inbox.post(message)
    setActive(inbox)
  }

  //注销独立节点 
  override def unregister(endpointName: String): Unit = synchronized {
    require(endpointName == name)
    //停止独立节点的inbox
    inbox.stop()
    // Mark active to handle the OnStop message.
    setActive(inbox)
    setActive(MessageLoop.PoisonPill)
    //线程池停止
    threadpool.shutdown()
  }
}

MessageLoop有一个基于单向链表实现的线程安全的阻塞队列active，该队列的每一个结点都是一个Inbox-收件箱，每个Inbox对应一个RpcEndpoint，有一个threadpool线程池，线程池里面的线程在启动之后都会执行receiveLoop方法，该方法每次都会取active队列的第一个结点-Inbox，如果该Inbox是毒丸PoisonPill则放回队列末尾停止处理消息，否则的话调用该Inbox的process方法，process方法就把消息发给该Inbox对应的RpcEndpoint端点，实际是调用RpcEndpoint端点对应的方法处理。

SharedMessageLoop继承自MessageLoop，提供Map类型变量-endpoints来保存RpcEndpoint端点和Inbox的映射关系，如果要给某个节点发消息，首先根据节点名在endpoints中找到对应的Inbox，然后调用该Inbox的post方法将消息发送过去，实际是放到Inbox内部的消息链表末尾等待处理，后续会讲到。共享消息循环在初始化的时候需要指定可用线程数，即指定有多少线程处理active队列中的消息。

总结：大部分场景用到的都是共享消息循环，其中最重要的是active队列以及线程池，线程池中的线程扫描队列，调用Inbox中的方法处理消息。