spark streaming 接收kafka消息之四 -- 运行在 worker 上的 receiver

使用分布式receiver来获取数据
使用 WAL 来实现 At least once 操作：
conf.set("spark.streaming.receiver.writeAheadLog.enable","true") // 开启 WAL
// 1、At most once - 每条数据最多被处理一次（0次或1次），这种语义下会出现数据丢失的问题；
// 2、At least once - 每条数据最少被处理一次 (1次或更多)，这个不会出现数据丢失，但是会出现数据重复；
// 3、Exactly once - 每条数据只会被处理一次，没有数据会丢失，并且没有数据会被多次处理，这种语义是大家最想要的，但是也是最难实现的。

如果不做容错，将会带来数据丢失，因为Receiver一直在接收数据，在其没有处理的时候（已通知zk数据接收到），Executor突然挂掉(或是driver挂掉通知executor关闭)，缓存在内存中的数据就会丢失。因为这个问题，Spark1.2开始加入了WAL（Write ahead log）开启 WAL，将receiver获取数据的存储级别修改为StorageLevel. MEMORY_AND_DISK_SER_2

// 缺点，不能自己维护消费 topic partition 的 offset
// 优点，开启 WAL，来确保 exactly-once 语义
val stream: ReceiverInputDStream[(String, String)] = KafkaUtils.createStream[String,String,StringDecoder,StringDecoder](
    ssc,kafkaParams,map,StorageLevel.MEMORY_AND_DISK_SER_2)

从Kafka 中读取数据

Driver 规划 receiver 运行的信息

org.apache.spark.streaming.StreamingContext#start中启动了 JobScheduler实例

// private[streaming] val scheduler = new JobScheduler(this)

// Start the streaming scheduler in a new thread, so that thread local properties
// like call sites and job groups can be reset without affecting those of the
// current thread.
ThreadUtils.runInNewThread("streaming-start") { // 单独的一个daemon线程运行函数题
  sparkContext.setCallSite(startSite.get)
  sparkContext.clearJobGroup()
  sparkContext.setLocalProperty(SparkContext.SPARK_JOB_INTERRUPT_ON_CANCEL, "false")
// 执行start 方法
  scheduler.start()
}
state = StreamingContextState.ACTIVE

org.apache.spark.streaming.scheduler.JobScheduler#start 源码如下：

def start(): Unit = synchronized {
  if (eventLoop != null) return // scheduler has already been started

  logDebug("Starting JobScheduler")
  eventLoop = new EventLoop[JobSchedulerEvent]("JobScheduler") {
    override protected def onReceive(event: JobSchedulerEvent): Unit = processEvent(event)

    override protected def onError(e: Throwable): Unit = reportError("Error in job scheduler", e)
  }
  eventLoop.start()

  // attach rate controllers of input streams to receive batch completion updates
  for {
    inputDStream <- ssc.graph.getInputStreams
    rateController <- inputDStream.rateController
  } ssc.addStreamingListener(rateController)

  listenerBus.start(ssc.sparkContext)
  receiverTracker = new ReceiverTracker(ssc)
  inputInfoTracker = new InputInfoTracker(ssc)
  receiverTracker.start()
  jobGenerator.start()
  logInfo("Started JobScheduler")
}

ReceiverTracker 的类声明如下：

This class manages the execution of the receivers of ReceiverInputDStreams. Instance of this class must be created after all input streams have been added and StreamingContext.start() has been called because it needs the final set of input streams at the time of instantiation.
此类负责执行ReceiverInputDStreams的receiver。必须在添加所有输入流并调用StreamingContext.start（）之后创建此类的实例，因为它在实例化时需要最终的输入流集。

其 start 方法如下：

/** Start the endpoint and receiver execution thread. */
def start(): Unit = synchronized {
  if (isTrackerStarted) {
    throw new SparkException("ReceiverTracker already started")
  }

  if (!receiverInputStreams.isEmpty) {
// 建立rpc endpoint
    endpoint = ssc.env.rpcEnv.setupEndpoint( // 注意，这是一个driver的 endpoint
      "ReceiverTracker", new ReceiverTrackerEndpoint(ssc.env.rpcEnv))
// driver节点上发送启动 receiver 命令
    if (!skipReceiverLaunch) launchReceivers()
    logInfo("ReceiverTracker started")
    trackerState = Started
  }
}

/**
 * Get the receivers from the ReceiverInputDStreams, distributes them to the
 * worker nodes as a parallel collection, and runs them.
 */
// 从ReceiverInputDStreams 获取到 receivers，然后将它们分配到不同的 worker 节点并运行它们。
private def launchReceivers(): Unit = {
  val receivers = receiverInputStreams.map(nis => {
// 未启用WAL 是KafkaReceiver，启动WAL后是ReliableKafkaReceiver
    val rcvr = nis.getReceiver()
    rcvr.setReceiverId(nis.id)
    rcvr
  })
  // 运行一个简单的应用来确保所有的salve node都已经启动起来，避免所有的 receiver 任务都在同一个local node上
  runDummySparkJob()

  logInfo("Starting " + receivers.length + " receivers")
  endpoint.send(StartAllReceivers(receivers)) // 发送请求driver 转发 启动 receiver 的命令
}

Driver 端StartAllReceivers 的处理代码如下：

override def receive: PartialFunction[Any, Unit] = {
  // Local messages
  case StartAllReceivers(receivers) =>
// schduleReceiver
    val scheduledLocations = schedulingPolicy.scheduleReceivers(receivers, getExecutors)
    for (receiver <- receivers) {
      val executors = scheduledLocations(receiver.streamId)
      updateReceiverScheduledExecutors(receiver.streamId, executors)
      receiverPreferredLocations(receiver.streamId) = receiver.preferredLocation
      startReceiver(receiver, executors)
    }
……
}

getExecutors源码如下：

/**
 * Get the list of executors excluding driver
 */
// 如果是 local 模式，返回 本地线程； 如果是 yarn 模式，返回 非driver 节点上的 excutors
private def getExecutors: Seq[ExecutorCacheTaskLocation] = {
  if (ssc.sc.isLocal) { // 如果在 local 模式下运行
    val blockManagerId = ssc.sparkContext.env.blockManager.blockManagerId
    Seq(ExecutorCacheTaskLocation(blockManagerId.host, blockManagerId.executorId))
  } else { // 在 yarn 模式下，过滤掉 driver 的 executor
    ssc.sparkContext.env.blockManager.master.getMemoryStatus.filter { case (blockManagerId, _) =>
      blockManagerId.executorId != SparkContext.DRIVER_IDENTIFIER // Ignore the driver location
    }.map { case (blockManagerId, _) =>
      ExecutorCacheTaskLocation(blockManagerId.host, blockManagerId.executorId)
    }.toSeq
  }
}

org.apache.spark.streaming.scheduler.ReceiverSchedulingPolicy#scheduleReceivers的解释如下：

Try our best to schedule receivers with evenly distributed. However, if the preferredLocations of receivers are not even, we may not be able to schedule them evenly because we have to respect them. Here is the approach to schedule executors:
First, schedule all the receivers with preferred locations (hosts), evenly among the executors running on those host.
Then, schedule all other receivers evenly among all the executors such that overall distribution over all the receivers is even.
This method is called when we start to launch receivers at the first time.
该方法就是确保receiver 能够在worker node 上均匀分布的。遵循以下两个原则：
1.使用 preferred location 分配 receiver 到这些node 上
2.将其他的未分配的receiver均匀分布均匀分布到 每一个 worker node 上 

org.apache.spark.streaming.scheduler.ReceiverTracker#updateReceiverScheduledExecutors 负责更新receiverid 和 receiver info 的映射关系，源码如下：

private def updateReceiverScheduledExecutors(
    receiverId: Int, scheduledLocations: Seq[TaskLocation]): Unit = {
  val newReceiverTrackingInfo = receiverTrackingInfos.get(receiverId) match {
    case Some(oldInfo) =>
      oldInfo.copy(state = ReceiverState.SCHEDULED,
        scheduledLocations = Some(scheduledLocations))
    case None =>
      ReceiverTrackingInfo(
        receiverId,
        ReceiverState.SCHEDULED,
        Some(scheduledLocations),
        runningExecutor = None)
  }
  receiverTrackingInfos.put(receiverId, newReceiverTrackingInfo)
}

Driver 发送分布式启动receiver job

startReceiver 负责启动 receiver，源码如下：

/**
 * Start a receiver along with its scheduled executors
 */
private def startReceiver(
    receiver: Receiver[_],
    scheduledLocations: Seq[TaskLocation]): Unit = {
  def shouldStartReceiver: Boolean = {
    // It's okay to start when trackerState is Initialized or Started
    !(isTrackerStopping || isTrackerStopped)
  }

  val receiverId = receiver.streamId
  if (!shouldStartReceiver) {
    onReceiverJobFinish(receiverId)
    return
  }

  val checkpointDirOption = Option(ssc.checkpointDir)
  val serializableHadoopConf =
    new SerializableConfiguration(ssc.sparkContext.hadoopConfiguration)

// 在 worker node 上启动 receiver 的方法
  val startReceiverFunc: Iterator[Receiver[_]] => Unit =
    (iterator: Iterator[Receiver[_]]) => {
      if (!iterator.hasNext) {
        throw new SparkException(
          "Could not start receiver as object not found.")
      }
      if (TaskContext.get().attemptNumber() == 0) {
        val receiver = iterator.next()
        assert(iterator.hasNext == false)
        val supervisor = new ReceiverSupervisorImpl(
          receiver, SparkEnv.get, serializableHadoopConf.value, checkpointDirOption)
        supervisor.start()
        supervisor.awaitTermination()
      } else {
        // It's restarted by TaskScheduler, but we want to reschedule it again. So exit it.
      }
    }

  // Create the RDD using the scheduledLocations to run the receiver in a Spark job
  val receiverRDD: RDD[Receiver[_]] =
    if (scheduledLocations.isEmpty) {
      ssc.sc.makeRDD(Seq(receiver), 1)
    } else {
      val preferredLocations = scheduledLocations.map(_.toString).distinct
      ssc.sc.makeRDD(Seq(receiver -> preferredLocations))
    }
  receiverRDD.setName(s"Receiver $receiverId")
  ssc.sparkContext.setJobDescription(s"Streaming job running receiver $receiverId")
  ssc.sparkContext.setCallSite(Option(ssc.getStartSite()).getOrElse(Utils.getCallSite()))
  // 提交分布式receiver 启动任务
  val future = ssc.sparkContext.submitJob[Receiver[_], Unit, Unit](
    receiverRDD, startReceiverFunc, Seq(0), (_, _) => Unit, ())
  // We will keep restarting the receiver job until ReceiverTracker is stopped
  future.onComplete {
    case Success(_) =>
      if (!shouldStartReceiver) {
        onReceiverJobFinish(receiverId)
      } else {
        logInfo(s"Restarting Receiver $receiverId")
        self.send(RestartReceiver(receiver))
      }
    case Failure(e) =>
      if (!shouldStartReceiver) {
        onReceiverJobFinish(receiverId)
      } else {
        logError("Receiver has been stopped. Try to restart it.", e)
        logInfo(s"Restarting Receiver $receiverId")
        self.send(RestartReceiver(receiver))
      }
  }(submitJobThreadPool)
  logInfo(s"Receiver ${receiver.streamId} started")
}

Worker节点启动 receiver监管服务

org.apache.spark.streaming.receiver.ReceiverSupervisorImpl#ReceiverSupervisorImpl 的 start 方法如下：

/** Start the supervisor */
def start() {
  onStart()
  startReceiver()
}
override protected def onStart() { // 启动 BlockGenerator 服务
  registeredBlockGenerators.foreach { _.start() }
}
// startReceiver 方法如下：
/** Start receiver */
def startReceiver(): Unit = synchronized {
  try {
    if (onReceiverStart()) { // 注册receiver 成功
      logInfo("Starting receiver")
      receiverState = Started
      receiver.onStart() // 启动 receiver
      logInfo("Called receiver onStart")
    } else {
      // The driver refused us
      stop("Registered unsuccessfully because Driver refused to start receiver " + streamId, None)
    }
  } catch {
    case NonFatal(t) =>
      stop("Error starting receiver " + streamId, Some(t))
  }
}

注册 receiver 到 driver节点

override protected def onReceiverStart(): Boolean = {
  val msg = RegisterReceiver(
    streamId, receiver.getClass.getSimpleName, host, executorId, endpoint)
  trackerEndpoint.askWithRetry[Boolean](msg)
}

简单描述一下driver 端做的事情，主要负责将其纳入到org.apache.spark.streaming.scheduler.ReceiverTracker 的管理中来，具体streamid 和 ReceiverTrackingInfo 的映射关系保存在receiverTrackingInfos中。

org.apache.spark.streaming.scheduler.ReceiverTracker#registerReceiver关键代码如下：

val name = s"${typ}-${streamId}"
val receiverTrackingInfo = ReceiverTrackingInfo(
  streamId,
  ReceiverState.ACTIVE,
  scheduledLocations = None,
  runningExecutor = Some(ExecutorCacheTaskLocation(host, executorId)),
  name = Some(name),
  endpoint = Some(receiverEndpoint))
receiverTrackingInfos.put(streamId, receiverTrackingInfo)
listenerBus.post(StreamingListenerReceiverStarted(receiverTrackingInfo.toReceiverInfo))

启动 receiver 线程

由于我们启用了 WAL， 所以 这里的receiver 是ReliableKafkaReceiver 的实例
receiver.onStart 即 org.apache.spark.streaming.kafka.ReliableKafkaReceiver#onStart, 源码如下：

override def onStart(): Unit = {
  logInfo(s"Starting Kafka Consumer Stream with group: $groupId")

  // Initialize the topic-partition / offset hash map.
// 1. 负责维护消费的 topic-partition 和 offset 的映射关系
  topicPartitionOffsetMap = new mutable.HashMap[TopicAndPartition, Long]

  // Initialize the stream block id / offset snapshot hash map.
// 2. 负责维护 block-id 和 partition-offset 之间的映射关系
  blockOffsetMap = new ConcurrentHashMap[StreamBlockId, Map[TopicAndPartition, Long]]()

  // Initialize the block generator for storing Kafka message.
// 3. 负责保存 kafka message 的 block generator，入参是GeneratedBlockHandler 实例，这是一个负责监听 block generator事件的一个监听器
// Generates batches of objects received by a org.apache.spark.streaming.receiver.Receiver and puts them into appropriately named blocks at regular intervals. This class starts two threads, one to periodically start a new batch and prepare the previous batch of as a block, the other to push the blocks into the block manager. 
  blockGenerator = supervisor.createBlockGenerator(new GeneratedBlockHandler)
  // 4. 关闭consumer 自动提交 offset 选项
// auto_offset_commit 应该是 false
  if (kafkaParams.contains(AUTO_OFFSET_COMMIT) && kafkaParams(AUTO_OFFSET_COMMIT) == "true") {
    logWarning(s"$AUTO_OFFSET_COMMIT should be set to false in ReliableKafkaReceiver, " +
      "otherwise we will manually set it to false to turn off auto offset commit in Kafka")
  }

  val props = new Properties()
  kafkaParams.foreach(param => props.put(param._1, param._2))
  // Manually set "auto.commit.enable" to "false" no matter user explicitly set it to true,
  // we have to make sure this property is set to false to turn off auto commit mechanism in Kafka.
  props.setProperty(AUTO_OFFSET_COMMIT, "false")

  val consumerConfig = new ConsumerConfig(props)

  assert(!consumerConfig.autoCommitEnable)

  logInfo(s"Connecting to Zookeeper: ${consumerConfig.zkConnect}")
// 5. 初始化 consumer 对象
// consumerConnector 是ZookeeperConsumerConnector的实例
  consumerConnector = Consumer.create(consumerConfig)
  logInfo(s"Connected to Zookeeper: ${consumerConfig.zkConnect}")
  // 6. 初始化zookeeper 的客户端
  zkClient = new ZkClient(consumerConfig.zkConnect, consumerConfig.zkSessionTimeoutMs,
    consumerConfig.zkConnectionTimeoutMs, ZKStringSerializer)
   // 7. 创建线程池来处理消息流，池的大小是固定的，为partition 的总数，并指定线程池中每一个线程的name 的前缀，内部使用ThreadPoolExecutor，并且 创建线程的 factory类是guava 工具包提供的。
  messageHandlerThreadPool = ThreadUtils.newDaemonFixedThreadPool(
    topics.values.sum, "KafkaMessageHandler")
   // 8. 启动 BlockGenerator内的两个线程
  blockGenerator.start()

// 9. 创建MessageStream对象
  val keyDecoder = classTag[U].runtimeClass.getConstructor(classOf[VerifiableProperties])
    .newInstance(consumerConfig.props)
    .asInstanceOf[Decoder[K]]

  val valueDecoder = classTag[T].runtimeClass.getConstructor(classOf[VerifiableProperties])
    .newInstance(consumerConfig.props)
    .asInstanceOf[Decoder[V]]
 
  val topicMessageStreams = consumerConnector.createMessageStreams(
    topics, keyDecoder, valueDecoder)
// 10. 将待处理的MessageHandler 放入 线程池中，等待执行
  topicMessageStreams.values.foreach { streams =>
    streams.foreach { stream =>
      messageHandlerThreadPool.submit(new MessageHandler(stream))
    }
  }
}

其中, 第9 步，创建MessageStream对象，
kafka.consumer.ZookeeperConsumerConnector#createMessageStreams 方法如下：

def createMessageStreams[K,V](topicCountMap: Map[String,Int], keyDecoder: Decoder[K], valueDecoder: Decoder[V])
    : Map[String, List[KafkaStream[K,V]]] = {
  if (messageStreamCreated.getAndSet(true))
    throw new MessageStreamsExistException(this.getClass.getSimpleName +
                                 " can create message streams at most once",null)
  consume(topicCountMap, keyDecoder, valueDecoder)
}

其调用了 consume 方法，源码如下：

def consume[K, V](topicCountMap: scala.collection.Map[String,Int], keyDecoder: Decoder[K], valueDecoder: Decoder[V])
    : Map[String,List[KafkaStream[K,V]]] = {
  debug("entering consume ")
  if (topicCountMap == null)
    throw new RuntimeException("topicCountMap is null")
 // 1. 初始化 topicCount
  val topicCount = TopicCount.constructTopicCount(consumerIdString, topicCountMap)
 // 2. 获取 每一个topic 和 threadId 集合的映射关系
  val topicThreadIds = topicCount.getConsumerThreadIdsPerTopic

  // make a list of (queue,stream) pairs, one pair for each threadId
// 3. 得到每一个 threadId 对应 (queue, stream) 的映射列表
  val queuesAndStreams = topicThreadIds.values.map(threadIdSet =>
    threadIdSet.map(_ => {
      val queue =  new LinkedBlockingQueue[FetchedDataChunk](config.queuedMaxMessages)
      val stream = new KafkaStream[K,V](
        queue, config.consumerTimeoutMs, keyDecoder, valueDecoder, config.clientId)
      (queue, stream)
    })
  ).flatten.toList
 // 4. 获取 groupId 在 zookeeper 中的path
  val dirs = new ZKGroupDirs(config.groupId)
// 5. 注册 consumer 到 groupId（在zk中）
  registerConsumerInZK(dirs, consumerIdString, topicCount)
// 6. 重新初始化 consumer
  reinitializeConsumer(topicCount, queuesAndStreams)
  // 7. 返回流 
  loadBalancerListener.kafkaMessageAndMetadataStreams.asInstanceOf[Map[String, List[KafkaStream[K,V]]]]
}

consumer消费kafka数据

在 kafka.consumer.ZookeeperConsumerConnector#consume方法中，有如下操作：

// 得到每一个 threadId 对应 (queue, stream) 的映射列表
  val queuesAndStreams = topicThreadIds.values.map(threadIdSet =>
    threadIdSet.map(_ => {
      val queue =  new LinkedBlockingQueue[FetchedDataChunk](config.queuedMaxMessages)
      val stream = new KafkaStream[K,V](
        queue, config.consumerTimeoutMs, keyDecoder, valueDecoder, config.clientId)
      (queue, stream)
    })
  ).flatten.toList
 // 获取 groupId 在 zookeeper 中的path
  val dirs = new ZKGroupDirs(config.groupId)
// 注册 consumer 到 groupId（在zk中）
  registerConsumerInZK(dirs, consumerIdString, topicCount)
// 重新初始化 consumer
  reinitializeConsumer(topicCount, queuesAndStreams)

在上面的代码中，可以看到初始化的queue(LinkedBlockingQueue实例)除了被传入stream(KafkaStream)的构造函数被迭代器从中取数据，还和 stream 重组成Tuple2[LinkedBlockingQueue[FetchedDataChunk]的list,之后被传入reinitializeConsumer 方法中。
kafka.consumer.ZookeeperConsumerConnector#reinitializeConsume 其源码如下：

private def reinitializeConsumer[K,V](
    topicCount: TopicCount,
    queuesAndStreams: List[(LinkedBlockingQueue[FetchedDataChunk],KafkaStream[K,V])]) {
 // 1. 获取 该groupid 在 zk 中的路径
  val dirs = new ZKGroupDirs(config.groupId)

  // listener to consumer and partition changes
// 2. 初始化loadBalancerListener，这个负载均衡listener 会时刻监控 consumer 和 partition 的变化
  if (loadBalancerListener == null) {
    val topicStreamsMap = new mutable.HashMap[String,List[KafkaStream[K,V]]]
    loadBalancerListener = new ZKRebalancerListener(
      config.groupId, consumerIdString, topicStreamsMap.asInstanceOf[scala.collection.mutable.Map[String, List[KafkaStream[_,_]]]])
  }

  // create listener for session expired event if not exist yet
  // 3. 监控 session 过期的listner， 有新session注册初始化，会通知 loadBalancer
if (sessionExpirationListener == null)
    sessionExpirationListener = new ZKSessionExpireListener(
      dirs, consumerIdString, topicCount, loadBalancerListener)

  // create listener for topic partition change event if not exist yet
// 4. 初始化ZKTopicPartitionChangeListener实例，当topic partition 变化时，这个listener会通知 loadBalancer
  if (topicPartitionChangeListener == null)
    topicPartitionChangeListener = new ZKTopicPartitionChangeListener(loadBalancerListener)
 // 5. 将queuesAndStreams 的值经过一系列转换，并添加到loadBalancerListener.kafkaMessageAndMetadataStreams 中
  val topicStreamsMap = loadBalancerListener.kafkaMessageAndMetadataStreams

  // map of {topic -> Set(thread-1, thread-2, ...)}
  val consumerThreadIdsPerTopic: Map[String, Set[ConsumerThreadId]] =
    topicCount.getConsumerThreadIdsPerTopic

  val allQueuesAndStreams = topicCount match {
    case wildTopicCount: WildcardTopicCount => // 这里是WildcardTopicCount，走这个分支
      /*
       * Wild-card consumption streams share the same queues, so we need to
       * duplicate the list for the subsequent zip operation.
       */
      (1 to consumerThreadIdsPerTopic.keySet.size).flatMap(_ => queuesAndStreams).toList
    case statTopicCount: StaticTopicCount =>
      queuesAndStreams
  }

  val topicThreadIds = consumerThreadIdsPerTopic.map {
    case(topic, threadIds) =>
      threadIds.map((topic, _))
  }.flatten

  require(topicThreadIds.size == allQueuesAndStreams.size,
    "Mismatch between thread ID count (%d) and queue count (%d)"
    .format(topicThreadIds.size, allQueuesAndStreams.size))
  val threadQueueStreamPairs = topicThreadIds.zip(allQueuesAndStreams)

  threadQueueStreamPairs.foreach(e => {
    val topicThreadId = e._1
    val q = e._2._1
    topicThreadIdAndQueues.put(topicThreadId, q)
    debug("Adding topicThreadId %s and queue %s to topicThreadIdAndQueues data structure".format(topicThreadId, q.toString))
    newGauge(
      "FetchQueueSize",
      new Gauge[Int] {
        def value = q.size
      },
      Map("clientId" -> config.clientId,
        "topic" -> topicThreadId._1,
        "threadId" -> topicThreadId._2.threadId.toString)
    )
  })

  val groupedByTopic = threadQueueStreamPairs.groupBy(_._1._1)
  groupedByTopic.foreach(e => {
    val topic = e._1
    val streams = e._2.map(_._2._2).toList
    topicStreamsMap += (topic -> streams)
    debug("adding topic %s and %d streams to map.".format(topic, streams.size))
  })

  // listener to consumer and partition changes
// 6. 使用 zkClient 注册sessionExpirationListener 实例
  zkClient.subscribeStateChanges(sessionExpirationListener)
 // 7. 使用 zkClient 注册loadBalancerListener 实例
  zkClient.subscribeChildChanges(dirs.consumerRegistryDir, loadBalancerListener)
 // 遍历每一个topic，使用zkClient 注册topicPartitionChangeListener 实例
  topicStreamsMap.foreach { topicAndStreams =>
    // register on broker partition path changes
    val topicPath = BrokerTopicsPath + "/" + topicAndStreams._1
    zkClient.subscribeDataChanges(topicPath, topicPartitionChangeListener)
  }

  // explicitly trigger load balancing for this consumer
// 8. 使用 loadBalancerListener 同步做负载均衡
  loadBalancerListener.syncedRebalance()
}

重点看 第 8 步，使用 loadBalancerListener 同步做负载均衡。
kafka.consumer.ZookeeperConsumerConnector.ZKRebalancerListener#syncedRebalance 源码如下：

def syncedRebalance() {
  rebalanceLock synchronized {
    rebalanceTimer.time {
      if(isShuttingDown.get())  { // 如果ZookeeperConsumerConnector
已经shutdown了，直接返回
        return
      } else {
        for (i <- 0 until config.rebalanceMaxRetries) { // 默认是 4 次
          info("begin rebalancing consumer " + consumerIdString + " try #" + i)
          var done = false
          var cluster: Cluster = null
          try {
            // 1. 根据zkClient 实例 获取并创建Cluster 对象，这个 cluster 实例包含了一个 Broker（broker的id，broker在zk中的路径） 列表
            cluster = getCluster(zkClient) 
            // 2. 在cluster中做 rebalance操作
            done = rebalance(cluster)
          } catch {
            case e: Throwable =>
              /** occasionally, we may hit a ZK exception because the ZK state is changing while we are iterating.
                * For example, a ZK node can disappear between the time we get all children and the time we try to get
                * the value of a child. Just let this go since another rebalance will be triggered.
                **/
              info("exception during rebalance ", e)
          }
          info("end rebalancing consumer " + consumerIdString + " try #" + i)
          if (done) {
            return
          } else {
            /* Here the cache is at a risk of being stale. To take future rebalancing decisions correctly, we should
             * clear the cache */
            info("Rebalancing attempt failed. Clearing the cache before the next rebalancing operation is triggered")
          }
          // stop all fetchers and clear all the queues to avoid data duplication
          closeFetchersForQueues(cluster, kafkaMessageAndMetadataStreams, topicThreadIdAndQueues.map(q => q._2))
          Thread.sleep(config.rebalanceBackoffMs)
        }
      }
    }
  }

  throw new ConsumerRebalanceFailedException(consumerIdString + " can't rebalance after " + config.rebalanceMaxRetries +" retries")
}

重点看 第2 步，在 cluster 中做 rebalance 操作，kafka.consumer.ZookeeperConsumerConnector.ZKRebalancerListener#rebalance 源码如下：

private def rebalance(cluster: Cluster): Boolean = {
  // 1. 获取 group和 threadId 的Map 映射关系
  val myTopicThreadIdsMap = TopicCount.constructTopicCount(
    group, consumerIdString, zkClient, config.excludeInternalTopics).getConsumerThreadIdsPerTopic
  // 2. 获取kafka cluster 中所有可用的node
  val brokers = getAllBrokersInCluster(zkClient)
  if (brokers.size == 0) { // 如果可用节点为空，设置listener订阅，返回 true
    // This can happen in a rare case when there are no brokers available in the cluster when the consumer is started.
    // We log an warning and register for child changes on brokers/id so that rebalance can be triggered when the brokers
    // are up.
    warn("no brokers found when trying to rebalance.")
    zkClient.subscribeChildChanges(ZkUtils.BrokerIdsPath, loadBalancerListener)
    true
  }
  else {
    /**
     * fetchers must be stopped to avoid data duplication, since if the current
     * rebalancing attempt fails, the partitions that are released could be owned by another consumer.
     * But if we don't stop the fetchers first, this consumer would continue returning data for released
     * partitions in parallel. So, not stopping the fetchers leads to duplicate data.
     */
   // 3. 做rebalance 之前的准备工作
   // 3.1. 关闭现有 fetcher 连接
    closeFetchers(cluster, kafkaMessageAndMetadataStreams, myTopicThreadIdsMap)
   // 3.2 释放 partition 的所有权（主要是删除zk下的owner 节点的数据以及解除内存中的topic和 fetcher的关联关系）
    releasePartitionOwnership(topicRegistry)
   // 3.3. 重新给partition分配 fetcher
    val assignmentContext = new AssignmentContext(group, consumerIdString, config.excludeInternalTopics, zkClient)
    val partitionOwnershipDecision = partitionAssignor.assign(assignmentContext)
    val currentTopicRegistry = new Pool[String, Pool[Int, PartitionTopicInfo]](
      valueFactory = Some((topic: String) => new Pool[Int, PartitionTopicInfo]))

    // fetch current offsets for all topic-partitions
    // 3.4 获取当前fetcher对应的 partitions 的 offsets，这里的offset是指 consumer 下一个要消费的offset
    val topicPartitions = partitionOwnershipDecision.keySet.toSeq

    val offsetFetchResponseOpt = fetchOffsets(topicPartitions)

    if (isShuttingDown.get || !offsetFetchResponseOpt.isDefined)
      false
    else {
      // 3.5 更新 partition 和 fetcher 的对应关系
      val offsetFetchResponse = offsetFetchResponseOpt.get
      topicPartitions.foreach(topicAndPartition => {
        val (topic, partition) = topicAndPartition.asTuple
// requestInfo是OffsetFetchResponse实例中的成员变量，它是一个Map[TopicAndPartition, OffsetMetadataAndError]实例
        val offset = offsetFetchResponse.requestInfo(topicAndPartition).offset
        val threadId = partitionOwnershipDecision(topicAndPartition)
        addPartitionTopicInfo(currentTopicRegistry, partition, topic, offset, threadId)
      })

      /**
       * move the partition ownership here, since that can be used to indicate a truly successful rebalancing attempt
       * A rebalancing attempt is completed successfully only after the fetchers have been started correctly
       */
      if(reflectPartitionOwnershipDecision(partitionOwnershipDecision)) {
        allTopicsOwnedPartitionsCount = partitionOwnershipDecision.size

        partitionOwnershipDecision.view.groupBy { case(topicPartition, consumerThreadId) => topicPartition.topic }
                                  .foreach { case (topic, partitionThreadPairs) =>
          newGauge("OwnedPartitionsCount",
            new Gauge[Int] {
              def value() = partitionThreadPairs.size
            },
            ownedPartitionsCountMetricTags(topic))
        }
        // 3.6 将已经新的 topic registry 覆盖旧的
        topicRegistry = currentTopicRegistry
// 4. 更新 fetcher
        updateFetcher(cluster)
        true
      } else {
        false
      }
    }
  }
}

其中addPartitionTopicInfo 源码如下：

private def addPartitionTopicInfo(currentTopicRegistry: Pool[String, Pool[Int, PartitionTopicInfo]],
                                    partition: Int, topic: String,
                                    offset: Long, consumerThreadId: ConsumerThreadId) {
//如果map没有对应的 key，会使用valueFactory初始化键值对，并返回 对应的 value
    val partTopicInfoMap = currentTopicRegistry.getAndMaybePut(topic)

    val queue = topicThreadIdAndQueues.get((topic, consumerThreadId))
    val consumedOffset = new AtomicLong(offset)
    val fetchedOffset = new AtomicLong(offset)
    val partTopicInfo = new PartitionTopicInfo(topic,
                                               partition,
                                               queue,
                                               consumedOffset,
                                               fetchedOffset,
                                               new AtomicInteger(config.fetchMessageMaxBytes),
                                               config.clientId)
    // 1. 将其注册到新的 Topic注册中心中，即注册 partition 和 fetcher 的关系
partTopicInfoMap.put(partition, partTopicInfo)
    debug(partTopicInfo + " selected new offset " + offset)
// 2. 更新consumer 的 已经消费的offset信息
    checkpointedZkOffsets.put(TopicAndPartition(topic, partition), offset)
  }
}

第4步， 更新 fetcher 源码如下：

private def updateFetcher(cluster: Cluster) {
  // update partitions for fetcher
  var allPartitionInfos : List[PartitionTopicInfo] = Nil
  for (partitionInfos <- topicRegistry.values)
    for (partition <- partitionInfos.values)
      allPartitionInfos ::= partition
  info("Consumer " + consumerIdString + " selected partitions : " +
    allPartitionInfos.sortWith((s,t) => s.partitionId < t.partitionId).map(_.toString).mkString(","))

  fetcher match {
    case Some(f) =>
      f.startConnections(allPartitionInfos, cluster)
    case None =>
  }
}

其中，f.startConnections方法真正执行 更新操作。此时引入一个新的类。即 fetcher 类，kafka.consumer.ConsumerFetcherManager。

kafka.consumer.ConsumerFetcherManager#startConnections 的源码如下：

def startConnections(topicInfos: Iterable[PartitionTopicInfo], cluster: Cluster) {
// LeaderFinderThread 在 topic 的leader node可用时，将 fetcher 添加到 leader 节点上
  leaderFinderThread = new LeaderFinderThread(consumerIdString + "-leader-finder-thread")
  leaderFinderThread.start()

  inLock(lock) {
// 更新ConsumerFetcherManager 成员变量
    partitionMap = topicInfos.map(tpi => (TopicAndPartition(tpi.topic, tpi.partitionId), tpi)).toMap
    this.cluster = cluster
    noLeaderPartitionSet ++= topicInfos.map(tpi => TopicAndPartition(tpi.topic, tpi.partitionId))
    cond.signalAll()
  }
}

ConsumerFetcherManager 有一个LeaderFinderThread 实例，该类的父类kafka.utils.ShutdownableThread ，run 方法如下：

override def run(): Unit = {
  info("Starting ")
  try{
    while(isRunning.get()){
      doWork()
    }
  } catch{
    case e: Throwable =>
      if(isRunning.get())
        error("Error due to ", e)
  }
  shutdownLatch.countDown()
  info("Stopped ")
}

doWork其实就是一个抽象方法，其子类LeaderFinderThread的实现如下：

// thread responsible for adding the fetcher to the right broker when leader is available
override def doWork() {
// 1. 获取 partition 和leader node的映射关系
  val leaderForPartitionsMap = new HashMap[TopicAndPartition, Broker]
  lock.lock()
  try {
    while (noLeaderPartitionSet.isEmpty) { // 这个字段在startConnections 已更新新值
      trace("No partition for leader election.")
      cond.await()
    }

    trace("Partitions without leader %s".format(noLeaderPartitionSet))
    val brokers = getAllBrokersInCluster(zkClient) // 获取所有可用broker 节点
    // 获取kafka.api.TopicMetadata 序列，kafka.api.TopicMetadata 保存了 topic 和 partitionId，isr，leader，replicas 的信息
val topicsMetadata = ClientUtils.fetchTopicMetadata(noLeaderPartitionSet.map(m => m.topic).toSet,
                                                        brokers,
                                                        config.clientId,
                                                        config.socketTimeoutMs,
                                                        correlationId.getAndIncrement).topicsMetadata
    if(logger.isDebugEnabled) topicsMetadata.foreach(topicMetadata => debug(topicMetadata.toString()))
// 2. 根据获取到的 partition 和 leader node 的关系更新noLeaderPartitionSet 和leaderForPartitionsMap 两个map集合，其中noLeaderPartitionSet 包含的是没有确定leader 的 partition 集合，leaderForPartitionsMap 是 已经确定了 leader 的 partition 集合。
    topicsMetadata.foreach { tmd =>
      val topic = tmd.topic
      tmd.partitionsMetadata.foreach { pmd =>
        val topicAndPartition = TopicAndPartition(topic, pmd.partitionId)
        if(pmd.leader.isDefined && noLeaderPartitionSet.contains(topicAndPartition)) {
          val leaderBroker = pmd.leader.get
          leaderForPartitionsMap.put(topicAndPartition, leaderBroker)
          noLeaderPartitionSet -= topicAndPartition
        }
      }
    }
  } catch {
    case t: Throwable => {
        if (!isRunning.get())
          throw t /* If this thread is stopped, propagate this exception to kill the thread. */
        else
          warn("Failed to find leader for %s".format(noLeaderPartitionSet), t)
      }
  } finally {
    lock.unlock()
  }

  try {
// 3. 具体为 partition 分配 fetcher
    addFetcherForPartitions(leaderForPartitionsMap.map{
      case (topicAndPartition, broker) =>
        topicAndPartition -> BrokerAndInitialOffset(broker, partitionMap(topicAndPartition).getFetchOffset())}
    )
  } catch {
    case t: Throwable => {
      if (!isRunning.get())
        throw t /* If this thread is stopped, propagate this exception to kill the thread. */
      else {
        warn("Failed to add leader for partitions %s; will retry".format(leaderForPartitionsMap.keySet.mkString(",")), t)
        lock.lock()
        noLeaderPartitionSet ++= leaderForPartitionsMap.keySet
        lock.unlock()
      }
    }
  }
  // 4. 关闭空闲fetcher线程
  shutdownIdleFetcherThreads()
  Thread.sleep(config.refreshLeaderBackoffMs)
}

重点看第3 步，具体为 partition 分配 fetcher，addFetcherForPartitions 源码如下：

def addFetcherForPartitions(partitionAndOffsets: Map[TopicAndPartition, BrokerAndInitialOffset]) {
  mapLock synchronized {
// 获取 fetcher 和 partition的映射关系
    val partitionsPerFetcher = partitionAndOffsets.groupBy{ case(topicAndPartition, brokerAndInitialOffset) =>
      BrokerAndFetcherId(brokerAndInitialOffset.broker, getFetcherId(topicAndPartition.topic, topicAndPartition.partition))}
    for ((brokerAndFetcherId, partitionAndOffsets) <- partitionsPerFetcher) {

      var fetcherThread: AbstractFetcherThread = null
      fetcherThreadMap.get(brokerAndFetcherId) match {
        case Some(f) => fetcherThread = f
        case None =>
// 根据brokerAndFetcherId 去初始化Fetcher并启动 fetcher
          fetcherThread = createFetcherThread(brokerAndFetcherId.fetcherId, brokerAndFetcherId.broker)
          fetcherThreadMap.put(brokerAndFetcherId, fetcherThread)
          fetcherThread.start
      }

      fetcherThreadMap(brokerAndFetcherId).addPartitions(partitionAndOffsets.map { case (topicAndPartition, brokerAndInitOffset) =>
        topicAndPartition -> brokerAndInitOffset.initOffset
      })
    }
  }

  info("Added fetcher for partitions %s".format(partitionAndOffsets.map{ case (topicAndPartition, brokerAndInitialOffset) =>
    "[" + topicAndPartition + ", initOffset " + brokerAndInitialOffset.initOffset + " to broker " + brokerAndInitialOffset.broker + "] "}))
}

kafka.consumer.ConsumerFetcherManager#createFetcherThread的源码如下：

override def createFetcherThread(fetcherId: Int, sourceBroker: Broker): AbstractFetcherThread = {
  new ConsumerFetcherThread(
    "ConsumerFetcherThread-%s-%d-%d".format(consumerIdString, fetcherId, sourceBroker.id),
    config, sourceBroker, partitionMap, this)
}

先来看ConsumerFetcherThread的构造方法声明：

class ConsumerFetcherThread(name: String,
                            val config: ConsumerConfig,
                            sourceBroker: Broker,
                            partitionMap: Map[TopicAndPartition, PartitionTopicInfo],
                            val consumerFetcherManager: ConsumerFetcherManager)
        extends AbstractFetcherThread(name = name, 
                                      clientId = config.clientId,
                                      sourceBroker = sourceBroker,
                                      socketTimeout = config.socketTimeoutMs,
                                      socketBufferSize = config.socketReceiveBufferBytes,
                                      fetchSize = config.fetchMessageMaxBytes,
                                      fetcherBrokerId = Request.OrdinaryConsumerId,
                                      maxWait = config.fetchWaitMaxMs,
                                      minBytes = config.fetchMinBytes,
                                      isInterruptible = true)

注意，partitionMap 中的value 是PartitionTopicInfo ，这个对象中封装了存放fetch结果值的BlockingQueue[FetchedDataChunk] 实例。
再来看 run 方法，其使用的是 kafka.utils.ShutdownableThread#run 方法，上面我们已经看过了，主要看该子类是如何重新 doWork方法的：

override def doWork() {
  inLock(partitionMapLock) { // 加锁，执行，释放锁
    if (partitionMap.isEmpty) // 如果没有需要执行的 fetch 操作，等待200ms后返回
      partitionMapCond.await(200L, TimeUnit.MILLISECONDS)
    partitionMap.foreach { // 将所有的 fetch 的信息添加到fetchRequestBuilder中
      case((topicAndPartition, offset)) =>
        fetchRequestBuilder.addFetch(topicAndPartition.topic, topicAndPartition.partition,
                         offset, fetchSize)
    }
  }
  // 构建批抓取的fetchRequest对象
  val fetchRequest = fetchRequestBuilder.build()
// 处理 FetchRequest
  if (!fetchRequest.requestInfo.isEmpty)
    processFetchRequest(fetchRequest)
}

其中 kafka.server.AbstractFetcherThread#processFetchRequest 源码如下：

private def processFetchRequest(fetchRequest: FetchRequest) {
  val partitionsWithError = new mutable.HashSet[TopicAndPartition]
  var response: FetchResponse = null
  try {
    trace("Issuing to broker %d of fetch request %s".format(sourceBroker.id, fetchRequest))
// 发送请求，并获取返回值。
// simpleConsumer  就是SimpleConsumer 实例，已作说明，不再赘述。
    response = simpleConsumer.fetch(fetchRequest)
  } catch {
    case t: Throwable =>
      if (isRunning.get) {
        warn("Error in fetch %s. Possible cause: %s".format(fetchRequest, t.toString))
        partitionMapLock synchronized {
          partitionsWithError ++= partitionMap.keys
        }
      }
  }
  fetcherStats.requestRate.mark()

  if (response != null) {
    // process fetched data
    inLock(partitionMapLock) { // 获取锁，执行处理response 操作，释放锁
      response.data.foreach {
        case(topicAndPartition, partitionData) =>
          val (topic, partitionId) = topicAndPartition.asTuple
          val currentOffset = partitionMap.get(topicAndPartition)
          // we append to the log if the current offset is defined and it is the same as the offset requested during fetch
          if (currentOffset.isDefined && fetchRequest.requestInfo(topicAndPartition).offset == currentOffset.get) {
            partitionData.error match { // 根据返回码来确定具体执行哪部分处理逻辑
              case ErrorMapping.NoError => // 成功返回，没有错误
                try {
                  val messages = partitionData.messages.asInstanceOf[ByteBufferMessageSet]
                  val validBytes = messages.validBytes
                  val newOffset = messages.shallowIterator.toSeq.lastOption match {
                    case Some(m: MessageAndOffset) => m.nextOffset
                    case None => currentOffset.get
                  }
                  partitionMap.put(topicAndPartition, newOffset)
                  fetcherLagStats.getFetcherLagStats(topic, partitionId).lag = partitionData.hw - newOffset
                  fetcherStats.byteRate.mark(validBytes)
                  // Once we hand off the partition data to the subclass, we can't mess with it any more in this thread
                  processPartitionData(topicAndPartition, currentOffset.get, partitionData)
                } catch {
                  case ime: InvalidMessageException => // 消息获取不完整
                    // we log the error and continue. This ensures two things
                    // 1. If there is a corrupt message in a topic partition, it does not bring the fetcher thread down and cause other topic partition to also lag
                    // 2. If the message is corrupt due to a transient state in the log (truncation, partial writes can cause this), we simply continue and
                    //    should get fixed in the subsequent fetches
                    logger.error("Found invalid messages during fetch for partition [" + topic + "," + partitionId + "] offset " + currentOffset.get + " error " + ime.getMessage)
                  case e: Throwable =>
                    throw new KafkaException("error processing data for partition [%s,%d] offset %d"
                                             .format(topic, partitionId, currentOffset.get), e)
                }
              case ErrorMapping.OffsetOutOfRangeCode => // offset out of range error
                try {
                  val newOffset = handleOffsetOutOfRange(topicAndPartition)
                  partitionMap.put(topicAndPartition, newOffset)
                  error("Current offset %d for partition [%s,%d] out of range; reset offset to %d"
                    .format(currentOffset.get, topic, partitionId, newOffset))
                } catch {
                  case e: Throwable =>
                    error("Error getting offset for partition [%s,%d] to broker %d".format(topic, partitionId, sourceBroker.id), e)
                    partitionsWithError += topicAndPartition
                }
              case _ =>
                if (isRunning.get) {
                  error("Error for partition [%s,%d] to broker %d:%s".format(topic, partitionId, sourceBroker.id,
                    ErrorMapping.exceptionFor(partitionData.error).getClass))
                  partitionsWithError += topicAndPartition
                }
            }
          }
      }
    }
  }

  if(partitionsWithError.size > 0) {
    debug("handling partitions with error for %s".format(partitionsWithError))
    handlePartitionsWithErrors(partitionsWithError)
  }
}

其中processPartitionData 源码如下，它负责处理具体的返回消息：

// process fetched data
def processPartitionData(topicAndPartition: TopicAndPartition, fetchOffset: Long, partitionData: FetchResponsePartitionData) {
// partitionMap 是一个成员变量，在构造函数中作为入参
  val pti = partitionMap(topicAndPartition)
  if (pti.getFetchOffset != fetchOffset)
    throw new RuntimeException("Offset doesn't match for partition [%s,%d] pti offset: %d fetch offset: %d"
                              .format(topicAndPartition.topic, topicAndPartition.partition, pti.getFetchOffset, fetchOffset))
// 数据入队
  pti.enqueue(partitionData.messages.asInstanceOf[ByteBufferMessageSet])
}

可以看到，终于在这里，把从leader中fetch的消息放入了BlockingQueue[FetchedDataChunk] 缓冲堵塞队列中。

KafkaStream从queue中堵塞式获取数据

KafkaStream 是依赖于 LinkedBlockingQueue 的同理 KafkaStream 也会返回一个迭代器 kafka.consumer.ConsumerIterator，用于迭代访问 KafkaStream 中的数据。
kafka.consumer.ConsumerIterator 的主要源码如下：

// 判断是否有下一个元素
def hasNext(): Boolean = {
  if(state == FAILED)
    throw new IllegalStateException("Iterator is in failed state")
  state match {
    case DONE => false
    case READY => true
    case _ => maybeComputeNext()
  }
}
// 获取下一个元素，父类实现
def next(): T = {
  if(!hasNext())
    throw new NoSuchElementException()
  state = NOT_READY
  if(nextItem == null)
    throw new IllegalStateException("Expected item but none found.")
  nextItem
}
// 获取下一个元素，使用子类ConsumerIterator实现
override def next(): MessageAndMetadata[K, V] = {
  val item = super.next() // 调用父类实现
  if(consumedOffset < 0)
    throw new KafkaException("Offset returned by the message set is invalid %d".format(consumedOffset))
  currentTopicInfo.resetConsumeOffset(consumedOffset)
  val topic = currentTopicInfo.topic
  trace("Setting %s consumed offset to %d".format(topic, consumedOffset))
  consumerTopicStats.getConsumerTopicStats(topic).messageRate.mark()
  consumerTopicStats.getConsumerAllTopicStats().messageRate.mark()
  item
}
 // 或许有，尝试计算一下下一个
def maybeComputeNext(): Boolean = {
  state = FAILED
  nextItem = makeNext()
  if(state == DONE) {
    false
  } else {
    state = READY
    true
  }
}
// 创建下一个元素，这个在子类ConsumerIterator中有实现
protected def makeNext(): MessageAndMetadata[K, V] = {
// 首先channel 是 LinkedBlockingQueue实例， 是 KafkaStream 中的 queue 成员变量，queue 成员变量
  var currentDataChunk: FetchedDataChunk = null
  // if we don't have an iterator, get one
  var localCurrent = current.get() 
// 如果没有迭代器或者是没有下一个元素了，需要从channel中取一个
  if(localCurrent == null || !localCurrent.hasNext) {
// 删除并返回队列的头节点。
    if (consumerTimeoutMs < 0)
      currentDataChunk = channel.take // 阻塞方法，一直等待，直到有可用元素
    else {
      currentDataChunk = channel.poll(consumerTimeoutMs,  TimeUnit.MILLISECONDS) // 阻塞方法，等待指定时间，超时也会返回
      if (currentDataChunk == null) { // 如果没有数据，重置状态为NOT_READY
        // reset state to make the iterator re-iterable
        resetState()
        throw new ConsumerTimeoutException
      }
    }
// 关闭命令
    if(currentDataChunk eq ZookeeperConsumerConnector.shutdownCommand) {
      debug("Received the shutdown command")
      return allDone // 该函数将状态设为DONE， 返回null
    } else {
      currentTopicInfo = currentDataChunk.topicInfo
      val cdcFetchOffset = currentDataChunk.fetchOffset
      val ctiConsumeOffset = currentTopicInfo.getConsumeOffset
      if (ctiConsumeOffset < cdcFetchOffset) {
        error("consumed offset: %d doesn't match fetch offset: %d for %s;
 Consumer may lose data"
          .format(ctiConsumeOffset, cdcFetchOffset, currentTopicInfo))
        currentTopicInfo.resetConsumeOffset(cdcFetchOffset)
      }
      localCurrent = currentDataChunk.messages.iterator

      current.set(localCurrent)
    }
    // if we just updated the current chunk and it is empty that means the fetch size is too small!
    if(currentDataChunk.messages.validBytes == 0)
      throw new MessageSizeTooLargeException("Found a message larger than the maximum fetch size of this consumer on topic " +
                                             "%s partition %d at fetch offset %d. Increase the fetch size, or decrease the maximum message size the broker will allow."
                                             .format(currentDataChunk.topicInfo.topic, currentDataChunk.topicInfo.partitionId, currentDataChunk.fetchOffset))
  }
  var item = localCurrent.next()
  // reject the messages that have already been consumed
  while (item.offset < currentTopicInfo.getConsumeOffset && localCurrent.hasNext) {
    item = localCurrent.next()
  }
  consumedOffset = item.nextOffset

  item.message.ensureValid() // validate checksum of message to ensure it is valid
 // 返回处理封装好的 kafka 数据
  new MessageAndMetadata(currentTopicInfo.topic, currentTopicInfo.partitionId, item.message, item.offset, keyDecoder, valueDecoder)
}

消费到的数据cache 到WAL中

我们再来看，org.apache.spark.streaming.kafka.ReliableKafkaReceiver#onStart 的第10 步相应的代码：

// 10. 将待处理的MessageHandler 放入 线程池中，等待执行
  topicMessageStreams.values.foreach { streams =>
    streams.foreach { stream =>
      messageHandlerThreadPool.submit(new MessageHandler(stream))
    }
  }

其中 MessageHandler 是一个 Runnable 对象，其 run 方法如下：

override def run(): Unit = {
  while (!isStopped) {
    try {
// 1. 获取ConsumerIterator 迭代器对象
      val streamIterator = stream.iterator()
      // 2. 遍历迭代器中获取每一条数据，并且保存message和相应的 metadata 信息
while (streamIterator.hasNext) {
        storeMessageAndMetadata(streamIterator.next)
      }
    } catch {
      case e: Exception =>
        reportError("Error handling message", e)
    }
  }
}

其中第二步中关键方法，org.apache.spark.streaming.kafka.ReliableKafkaReceiver#storeMessageAndMetadata 方法如下：

/** Store a Kafka message and the associated metadata as a tuple. */
private def storeMessageAndMetadata(
    msgAndMetadata: MessageAndMetadata[K, V]): Unit = {
  val topicAndPartition = TopicAndPartition(msgAndMetadata.topic, msgAndMetadata.partition)
  val data = (msgAndMetadata.key, msgAndMetadata.message)
  val metadata = (topicAndPartition, msgAndMetadata.offset)
// 添加数据到 block
  blockGenerator.addDataWithCallback(data, metadata)
}

addDataWithCallback 源码如下：

/**
 * Push a single data item into the buffer. After buffering the data, the
 * `BlockGeneratorListener.onAddData` callback will be called.
 */
def addDataWithCallback(data: Any, metadata: Any): Unit = {
  if (state == Active) {
    waitToPush()
    synchronized {
      if (state == Active) {
// 1. 将数据放入 buffer 中，以便处理线程从中获取数据
        currentBuffer += data
// 2. 在启动 receiver线程中，可以知道listener 是指GeneratedBlockHandler 实例
        listener.onAddData(data, metadata)
      } else {
        throw new SparkException(
          "Cannot add data as BlockGenerator has not been started or has been stopped")
      }
    }
  } else {
    throw new SparkException(
      "Cannot add data as BlockGenerator has not been started or has been stopped")
  }
}

第二步比较简单，先看一下第二步：
org.apache.spark.streaming.kafka.ReliableKafkaReceiver.GeneratedBlockHandler#onAddData的源码如下：

def onAddData(data: Any, metadata: Any): Unit = {
  // Update the offset of the data that was added to the generator
  if (metadata != null) {
    val (topicAndPartition, offset) = metadata.asInstanceOf[(TopicAndPartition, Long)]
    updateOffset(topicAndPartition, offset)
  }
}
// 这里的 updateOffset 调用的是//org.apache.spark.streaming.kafka.ReliableKafkaReceiver#updateOffset，源码如下：
/** Update stored offset */
private def updateOffset(topicAndPartition: TopicAndPartition, offset: Long): Unit = {
  topicPartitionOffsetMap.put(topicAndPartition, offset)
}

第一步的原理如下：
在 BlockGenerator中有一个定时器，定时（200ms）去执行检查currentBuffer是否为empty任务， 若不为空，则执行如下操作并把它放入等待生成block 的队列中，有两外一个线程来时刻监听这个队列，有数据，则执行pushBlock 操作。
第一个定时器线程如下：

private val blockIntervalTimer =
  new RecurringTimer(clock, blockIntervalMs, updateCurrentBuffer, "BlockGenerator")

// 其中，updateCurrentBuffer 方法如下
/** Change the buffer to which single records are added to. */
private def updateCurrentBuffer(time: Long): Unit = {
  try {
    var newBlock: Block = null
    synchronized {
      if (currentBuffer.nonEmpty) {
        val newBlockBuffer = currentBuffer
        currentBuffer = new ArrayBuffer[Any]
        val blockId = StreamBlockId(receiverId, time - blockIntervalMs)
        listener.onGenerateBlock(blockId)
        newBlock = new Block(blockId, newBlockBuffer)
      }
    }

    if (newBlock != null) {
      blocksForPushing.put(newBlock)  // put is blocking when queue is full
    }
  } catch {
    case ie: InterruptedException =>
      logInfo("Block updating timer thread was interrupted")
    case e: Exception =>
      reportError("Error in block updating thread", e)
  }
}

// listener.onGenerateBlock(blockId) 代码如下：
def onGenerateBlock(blockId: StreamBlockId): Unit = {
  // Remember the offsets of topics/partitions when a block has been generated
  rememberBlockOffsets(blockId)
}
// rememberBlockOffsets 代码如下：
/**
 * Remember the current offsets for each topic and partition. This is called when a block is
 * generated.
 */
private def rememberBlockOffsets(blockId: StreamBlockId): Unit = {
  // Get a snapshot of current offset map and store with related block id.
  val offsetSnapshot = topicPartitionOffsetMap.toMap
  blockOffsetMap.put(blockId, offsetSnapshot)
  topicPartitionOffsetMap.clear()
}
// 可以看出，主要是清除 topic-partition-> offset 映射关系
// 建立 block 和topic-partition-> offset的映射关系

其中，blocksForPushing是一个有界阻塞队列，另外一个线程会一直轮询它。

private val blocksForPushing = new ArrayBlockingQueue[Block](blockQueueSize)
private val blockPushingThread = new Thread() { override def run() { keepPushingBlocks() } }

/** Keep pushing blocks to the BlockManager. */
// 这个方法主要的作用就是一直不停地轮询blocksForPushing队列，并处理相应的push block 事件。
private def keepPushingBlocks() {
  logInfo("Started block pushing thread")

  def areBlocksBeingGenerated: Boolean = synchronized {
    state != StoppedGeneratingBlocks
  }

  try {
    // While blocks are being generated, keep polling for to-be-pushed blocks and push them.
    while (areBlocksBeingGenerated) { // 线程没有被停止，则一直循环
// 超时poll操作获取并删除头节点，超过时间（10ms）则返回
      Option(blocksForPushing.poll(10, TimeUnit.MILLISECONDS)) match {
        case Some(block) => pushBlock(block) // 如果有数据则进行处理。
        case None =>
      }
    }

    // At this point, state is StoppedGeneratingBlock. So drain the queue of to-be-pushed blocks.
    logInfo("Pushing out the last " + blocksForPushing.size() + " blocks")
    while (!blocksForPushing.isEmpty) { // 如果队列中还有数据，继续进行处理
      val block = blocksForPushing.take() // 这是一个堵塞方法，不过现在会马上返回，因为队列里面有数据。
      logDebug(s"Pushing block $block")
      pushBlock(block) // 处理数据
      logInfo("Blocks left to push " + blocksForPushing.size())
    }
    logInfo("Stopped block pushing thread")
  } catch {
    case ie: InterruptedException =>
      logInfo("Block pushing thread was interrupted")
    case e: Exception =>
      reportError("Error in block pushing thread", e)
  }
}

其中的pushBlock源码如下：

private def pushBlock(block: Block) {
  listener.onPushBlock(block.id, block.buffer)
  logInfo("Pushed block " + block.id)
}

其调用的listener(org.apache.spark.streaming.kafka.ReliableKafkaReceiver.GeneratedBlockHandler)的 onPushBlock 源码如下：

def onPushBlock(blockId: StreamBlockId, arrayBuffer: mutable.ArrayBuffer[_]): Unit = {
  // Store block and commit the blocks offset
  storeBlockAndCommitOffset(blockId, arrayBuffer)
}

其中，storeBlockAndCommitOffset具体代码如下：

/**
 * Store the ready-to-be-stored block and commit the related offsets to zookeeper. This method
 * will try a fixed number of times to push the block. If the push fails, the receiver is stopped.
 */
private def storeBlockAndCommitOffset(
    blockId: StreamBlockId, arrayBuffer: mutable.ArrayBuffer[_]): Unit = {
  var count = 0
  var pushed = false
  var exception: Exception = null
  while (!pushed && count <= 3) { // 整个过程，总共允许3 次重试
    try {
      store(arrayBuffer.asInstanceOf[mutable.ArrayBuffer[(K, V)]])
      pushed = true
    } catch {
      case ex: Exception =>
        count += 1
        exception = ex
    }
  }
  if (pushed) { // 已经push block
// 更新 offset
    Option(blockOffsetMap.get(blockId)).foreach(commitOffset)
// 如果已经push 到 BlockManager 中，则不会再保留 block和topic-partition-> offset的映射关系
    blockOffsetMap.remove(blockId)
  } else {
    stop("Error while storing block into Spark", exception)
  }
}
// 其中，commitOffset源码如下：
/**
 * Commit the offset of Kafka's topic/partition, the commit mechanism follow Kafka 0.8.x's
 * metadata schema in Zookeeper.
 */
private def commitOffset(offsetMap: Map[TopicAndPartition, Long]): Unit = {
  if (zkClient == null) {
    val thrown = new IllegalStateException("Zookeeper client is unexpectedly null")
    stop("Zookeeper client is not initialized before commit offsets to ZK", thrown)
    return
  }

  for ((topicAndPart, offset) <- offsetMap) {
    try {
// 获取在 zk 中 comsumer 的partition的目录
      val topicDirs = new ZKGroupTopicDirs(groupId, topicAndPart.topic)
      val zkPath = s"${topicDirs.consumerOffsetDir}/${topicAndPart.partition}"
      // 更新 consumer 的已消费topic-partition 的offset 操作
      ZkUtils.updatePersistentPath(zkClient, zkPath, offset.toString)
    } catch {
      case e: Exception =>
        logWarning(s"Exception during commit offset $offset for topic" +
          s"${topicAndPart.topic}, partition ${topicAndPart.partition}", e)
    }

    logInfo(s"Committed offset $offset for topic ${topicAndPart.topic}, " +
      s"partition ${topicAndPart.partition}")
  }
}

关键方法store 如下：

/** Store an ArrayBuffer of received data as a data block into Spark's memory. */
def store(dataBuffer: ArrayBuffer[T]) {
  supervisor.pushArrayBuffer(dataBuffer, None, None)
}

其调用了supervisor（org.apache.spark.streaming.receiver.ReceiverSupervisorImpl实例）的pushArrayBuffer方法，内部操作如下：

/** Store an ArrayBuffer of received data as a data block into Spark's memory. */
def pushArrayBuffer(
    arrayBuffer: ArrayBuffer[_],
    metadataOption: Option[Any],
    blockIdOption: Option[StreamBlockId]
  ) {
  pushAndReportBlock(ArrayBufferBlock(arrayBuffer), metadataOption, blockIdOption)
}

org.apache.spark.streaming.receiver.ReceiverSupervisorImpl#pushAndReportBlock 源码如下：

/** Store block and report it to driver */
def pushAndReportBlock(
    receivedBlock: ReceivedBlock,
    metadataOption: Option[Any],
    blockIdOption: Option[StreamBlockId]
  ) {
// 1.准备blockId，time等信息
  val blockId = blockIdOption.getOrElse(nextBlockId)
  val time = System.currentTimeMillis
// 2. 执行存储 block 操作
  val blockStoreResult = receivedBlockHandler.storeBlock(blockId, receivedBlock)
  logDebug(s"Pushed block $blockId in ${(System.currentTimeMillis - time)} ms")
// 3. 获取保存的message 的记录数
  val numRecords = blockStoreResult.numRecords
// 4. 通知trackerEndpoint已经添加block，执行更新driver 的WAL操作
  val blockInfo = ReceivedBlockInfo(streamId, numRecords, metadataOption, blockStoreResult)
  trackerEndpoint.askWithRetry[Boolean](AddBlock(blockInfo))
  logDebug(s"Reported block $blockId")
}

其中，receivedBlockHandler 的赋值语句如下：

private val receivedBlockHandler: ReceivedBlockHandler = {
  if (WriteAheadLogUtils.enableReceiverLog(env.conf)) {
    if (checkpointDirOption.isEmpty) {
      throw new SparkException(
        "Cannot enable receiver write-ahead log without checkpoint directory set. " +
          "Please use streamingContext.checkpoint() to set the checkpoint directory. " +
          "See documentation for more details.")
    }
// enable WAL并且checkpoint dir 不为空，即，在这里，返回WriteAheadLogBasedBlockHandler 对象，这个对象持有了 blockmanager，streamid，storagelevel，conf，checkpointdir 等信息
    new WriteAheadLogBasedBlockHandler(env.blockManager, receiver.streamId,
      receiver.storageLevel, env.conf, hadoopConf, checkpointDirOption.get)
  } else {
    new BlockManagerBasedBlockHandler(env.blockManager, receiver.storageLevel)
  }
}

ReceivedBlockHandler 的 storeBlock方法源码如下：

/**
 * This implementation stores the block into the block manager as well as a write ahead log.
 * It does this in parallel, using Scala Futures, and returns only after the block has
 * been stored in both places.
 */
// 并行地将block 存入 blockmanager 和 write ahead log，使用scala 的Future 机制实现的，当两个都写完毕之后，返回。
def storeBlock(blockId: StreamBlockId, block: ReceivedBlock): ReceivedBlockStoreResult = {

  var numRecords = None: Option[Long]
  // Serialize the block so that it can be inserted into both
// 1. 将ReceivedBlock序列化（未使用压缩机制）成字节数组
  val serializedBlock = block match { // serializedBlock 就是序列化后的结果
    case ArrayBufferBlock(arrayBuffer) => // go this branch
      numRecords = Some(arrayBuffer.size.toLong)
      blockManager.dataSerialize(blockId, arrayBuffer.iterator)
    case IteratorBlock(iterator) =>
      val countIterator = new CountingIterator(iterator)
      val serializedBlock = blockManager.dataSerialize(blockId, countIterator)
      numRecords = countIterator.count
      serializedBlock
    case ByteBufferBlock(byteBuffer) =>
      byteBuffer
    case _ =>
      throw new Exception(s"Could not push $blockId to block manager, unexpected block type")
  }

  // 2. Store the block in block manager
  val storeInBlockManagerFuture = Future {
    val putResult =
      blockManager.putBytes(blockId, serializedBlock, effectiveStorageLevel, tellMaster = true)
    if (!putResult.map { _._1 }.contains(blockId)) {
      throw new SparkException(
        s"Could not store $blockId to block manager with storage level $storageLevel")
    }
  }

  // 3. Store the block in write ahead log
  val storeInWriteAheadLogFuture = Future {
    writeAheadLog.write(serializedBlock, clock.getTimeMillis())
  }

  // 4. Combine the futures, wait for both to complete, and return the write ahead log record handle
  val combinedFuture = storeInBlockManagerFuture.zip(storeInWriteAheadLogFuture).map(_._2)
// 等待future任务结果返回。默认时间是 30s， 使用spark.streaming.receiver.blockStoreTimeout 参数来变更默认值
  val walRecordHandle = Await.result(combinedFuture, blockStoreTimeout)
  // 返回cache之后的block 相关信息
WriteAheadLogBasedStoreResult(blockId, numRecords, walRecordHandle)
}

将WAL的block信息发送给driver

注意WriteAheadLogBasedStoreResult 这个 WriteAheadLogBasedStoreResult 实例，后面 RDD 在处理的时候会使用到。
org.apache.spark.streaming.receiver.ReceiverSupervisorImpl#pushAndReportBlock 通知driver addBlock 的源码如下：

// 4. 通知trackerEndpoint已经添加block，执行更新driver 的WAL操作
  val blockInfo = ReceivedBlockInfo(streamId, numRecords, metadataOption, blockStoreResult)
  trackerEndpoint.askWithRetry[Boolean](AddBlock(blockInfo))
  logDebug(s"Reported block $blockId")

Driver将WAL block数据写入到 driver 的WAL中

跳过中间的RPC操作，直接到 driver 端org.apache.spark.streaming.scheduler.ReceiverTracker.ReceiverTrackerEndpoint#receiveAndReply 中：

case AddBlock(receivedBlockInfo) =>
  if (WriteAheadLogUtils.isBatchingEnabled(ssc.conf, isDriver = true)) {
    walBatchingThreadPool.execute(new Runnable {
      override def run(): Unit = Utils.tryLogNonFatalError {
        if (active) {
          context.reply(addBlock(receivedBlockInfo))
        } else {
          throw new IllegalStateException("ReceiverTracker RpcEndpoint shut down.")
        }
      }
    })
  } else {
    context.reply(addBlock(receivedBlockInfo))
  }

其中 addBlock方法源码如下：

/** Add new blocks for the given stream */
private def addBlock(receivedBlockInfo: ReceivedBlockInfo): Boolean = {
  receivedBlockTracker.addBlock(receivedBlockInfo)
}

其中，org.apache.spark.streaming.scheduler.ReceivedBlockTracker#addBlock 源码如下：

/** Add received block. This event will get written to the write ahead log (if enabled). */
def addBlock(receivedBlockInfo: ReceivedBlockInfo): Boolean = {
  try {
    val writeResult = writeToLog(BlockAdditionEvent(receivedBlockInfo))
    if (writeResult) {
      synchronized {
        getReceivedBlockQueue(receivedBlockInfo.streamId) += receivedBlockInfo
      }
      logDebug(s"Stream ${receivedBlockInfo.streamId} received " +
        s"block ${receivedBlockInfo.blockStoreResult.blockId}")
    } else {
      logDebug(s"Failed to acknowledge stream ${receivedBlockInfo.streamId} receiving " +
        s"block ${receivedBlockInfo.blockStoreResult.blockId} in the Write Ahead Log.")
    }
    writeResult
  } catch {
    case NonFatal(e) =>
      logError(s"Error adding block $receivedBlockInfo", e)
      false
  }
}
/** Write an update to the tracker to the write ahead log */
private def writeToLog(record: ReceivedBlockTrackerLogEvent): Boolean = {
  if (isWriteAheadLogEnabled) {
    logTrace(s"Writing record: $record")
    try {
      writeAheadLogOption.get.write(ByteBuffer.wrap(Utils.serialize(record)),
        clock.getTimeMillis())
      true
    } catch {
      case NonFatal(e) =>
        logWarning(s"Exception thrown while writing record: $record to the WriteAheadLog.", e)
        false
    }
  } else {
    true
  }
}
/** Get the queue of received blocks belonging to a particular stream */
private def getReceivedBlockQueue(streamId: Int): ReceivedBlockQueue = {
  streamIdToUnallocatedBlockQueues.getOrElseUpdate(streamId, new ReceivedBlockQueue)
}

上述代码，主要是将BlockAdditionEvent写WAL和更新队列（其实就是mutable.HashMap[Int, ReceivedBlockQueue]），这个队列中存放的是streamId ->UnallocatedBlock 的映射关系

从WAL RDD 中读取数据

createStream 源码如下：

/**
 * Create an input stream that pulls messages from Kafka Brokers.
 * @param ssc         StreamingContext object
 * @param kafkaParams Map of kafka configuration parameters,
 *                    see http://kafka.apache.org/08/configuration.html
 * @param topics      Map of (topic_name -> numPartitions) to consume. Each partition is consumed
 *                    in its own thread.
 * @param storageLevel Storage level to use for storing the received objects
 * @tparam K type of Kafka message key
 * @tparam V type of Kafka message value
 * @tparam U type of Kafka message key decoder
 * @tparam T type of Kafka message value decoder
 * @return DStream of (Kafka message key, Kafka message value)
 */
def createStream[K: ClassTag, V: ClassTag, U <: Decoder[_]: ClassTag, T <: Decoder[_]: ClassTag](
    ssc: StreamingContext,
    kafkaParams: Map[String, String],
    topics: Map[String, Int],
    storageLevel: StorageLevel
  ): ReceiverInputDStream[(K, V)] = {
// 可以通过设置spark.streaming.receiver.writeAheadLog.enable参数为 true来开启WAL
  val walEnabled = WriteAheadLogUtils.enableReceiverLog(ssc.conf)
  new KafkaInputDStream[K, V, U, T](ssc, kafkaParams, topics, walEnabled, storageLevel)
}

创建的是KafkaInputDStream对象：

/**
 * Input stream that pulls messages from a Kafka Broker.
 *
 * @param kafkaParams Map of kafka configuration parameters.
 *                    See: http://kafka.apache.org/configuration.html
 * @param topics Map of (topic_name -> numPartitions) to consume. Each partition is consumed
 * in its own thread.
 * @param storageLevel RDD storage level.
 */
private[streaming]
class KafkaInputDStream[
  K: ClassTag,
  V: ClassTag,
  U <: Decoder[_]: ClassTag,
  T <: Decoder[_]: ClassTag](
    ssc_ : StreamingContext,
    kafkaParams: Map[String, String],
    topics: Map[String, Int],
    useReliableReceiver: Boolean,
    storageLevel: StorageLevel
  ) extends ReceiverInputDStream[(K, V)](ssc_) with Logging {

  def getReceiver(): Receiver[(K, V)] = {
    if (!useReliableReceiver) { // 未启用 WAL，会使用 KafkaReceiver 对象
      new KafkaReceiver[K, V, U, T](kafkaParams, topics, storageLevel)
    } else { // 如果启用了WAL， 使用ReliableKafkaReceiver
      new ReliableKafkaReceiver[K, V, U, T](kafkaParams, topics, storageLevel)
    }
  }
}

org.apache.spark.streaming.kafka.KafkaInputDStream 继承父类的 compute方法：

/**
 * Generates RDDs with blocks received by the receiver of this stream. */
override def compute(validTime: Time): Option[RDD[T]] = {
  val blockRDD = {

    if (validTime < graph.startTime) {
      // If this is called for any time before the start time of the context,
      // then this returns an empty RDD. This may happen when recovering from a
      // driver failure without any write ahead log to recover pre-failure data.
      new BlockRDD[T](ssc.sc, Array.empty)
    } else {
      // Otherwise, ask the tracker for all the blocks that have been allocated to this stream
      // for this batch
      val receiverTracker = ssc.scheduler.receiverTracker
      val blockInfos = receiverTracker.getBlocksOfBatch(validTime).getOrElse(id, Seq.empty)

      // Register the input blocks information into InputInfoTracker
      val inputInfo = StreamInputInfo(id, blockInfos.flatMap(_.numRecords).sum)
      ssc.scheduler.inputInfoTracker.reportInfo(validTime, inputInfo)

      // Create the BlockRDD
      createBlockRDD(validTime, blockInfos)
    }
  }
  Some(blockRDD)
}

getBlocksOfBatch 如下：

/** Get the blocks for the given batch and all input streams. */
def getBlocksOfBatch(batchTime: Time): Map[Int, Seq[ReceivedBlockInfo]] = {
  receivedBlockTracker.getBlocksOfBatch(batchTime)
}
调用：
/** Get the blocks allocated to the given batch. */
def getBlocksOfBatch(batchTime: Time): Map[Int, Seq[ReceivedBlockInfo]] = synchronized {
  timeToAllocatedBlocks.get(batchTime).map { _.streamIdToAllocatedBlocks }.getOrElse(Map.empty)
}

JobGenerator将WAL block 分配给一个batch，并生成job

取出WAL block 信息

在 org.apache.spark.streaming.scheduler.JobGenerator 中声明了一个定时器：

// timer 会按照批次间隔 生成 GenerateJobs 任务，并放入eventLoop 堵塞队列中
private val timer = new RecurringTimer(clock, ssc.graph.batchDuration.milliseconds,
  longTime => eventLoop.post(GenerateJobs(new Time(longTime))), "JobGenerator")

EventLoop 实例化代码如下：

eventLoop = new EventLoop[JobGeneratorEvent]("JobGenerator") {
  override protected def onReceive(event: JobGeneratorEvent): Unit = processEvent(event)

  override protected def onError(e: Throwable): Unit = {
    jobScheduler.reportError("Error in job generator", e)
  }
}
eventLoop.start()

EventLoop里定义了一个LinkedBlockingDeque双端堵塞队列和一个执行daemon线程，daemon线程会不停从 双端堵塞队列中堵塞式取数据，一旦取到数据，会调 onReceive 方法，即 processEvent 方法：

/** Processes all events */
private def processEvent(event: JobGeneratorEvent) {
  logDebug("Got event " + event)
  event match {
    case GenerateJobs(time) => generateJobs(time)
    case ClearMetadata(time) => clearMetadata(time)
    case DoCheckpoint(time, clearCheckpointDataLater) =>
      doCheckpoint(time, clearCheckpointDataLater)
    case ClearCheckpointData(time) => clearCheckpointData(time)
  }
}

由于是GenerateJobs 事件， 会继续调用generateJobs 方法：

/** Generate jobs and perform checkpoint for the given `time`.  */
private def generateJobs(time: Time) {
  // Set the SparkEnv in this thread, so that job generation code can access the environment
  // Example: BlockRDDs are created in this thread, and it needs to access BlockManager
  // Update: This is probably redundant after threadlocal stuff in SparkEnv has been removed.
  SparkEnv.set(ssc.env)
  Try {
// 1. 将 WAL block 信息 分配给batch（这些数据块信息是worker 节点cache 到WAL 之后发送给driver 端的）
    jobScheduler.receiverTracker.allocateBlocksToBatch(time)
// 2. 使用分配的block数据块来生成任务
    graph.generateJobs(time) // generate jobs using allocated block
  } match {
    case Success(jobs) =>
      val streamIdToInputInfos = jobScheduler.inputInfoTracker.getInfo(time)
      jobScheduler.submitJobSet(JobSet(time, jobs, streamIdToInputInfos))
    case Failure(e) =>
      jobScheduler.reportError("Error generating jobs for time " + time, e)
  }
// 发布DoCheckpoint 事件，保存checkpoint操作，主要是将新的checkpoint 数据写入到 hdfs， 删除旧的 checkpoint 数据
  eventLoop.post(DoCheckpoint(time, clearCheckpointDataLater = false))
}

第一步中调用的
org.apache.spark.streaming.scheduler.ReceiverTracker#allocateBlocksToBatch方法如下：

/** Allocate all unallocated blocks to the given batch. */
def allocateBlocksToBatch(batchTime: Time): Unit = {
  if (receiverInputStreams.nonEmpty) {
    receivedBlockTracker.allocateBlocksToBatch(batchTime)
  }
}

其中，org.apache.spark.streaming.scheduler.ReceivedBlockTracker#allocateBlocksToBatch 方法如下：

def allocateBlocksToBatch(batchTime: Time): Unit = synchronized {
  if (lastAllocatedBatchTime == null || batchTime > lastAllocatedBatchTime) {
// 遍历输入流，根据流的 streamId 获取未被分配的block队列，并返回[streamId, seq[receivedBlockInfo]],由此可知，到此为止，数据其实已经从receiver中读出来了。
   // 获取 streamid和 WAL的blocks 的映射关系
val streamIdToBlocks = streamIds.map { streamId =>
        (streamId, getReceivedBlockQueue(streamId).dequeueAll(x => true))
    }.toMap
    val allocatedBlocks = AllocatedBlocks(streamIdToBlocks)
    if (writeToLog(BatchAllocationEvent(batchTime, allocatedBlocks))) {
      timeToAllocatedBlocks.put(batchTime, allocatedBlocks)
      lastAllocatedBatchTime = batchTime
    } else {
      logInfo(s"Possibly processed batch $batchTime need to be processed again in WAL recovery")
    }
  } else {
    // This situation occurs when:
    // 1. WAL is ended with BatchAllocationEvent, but without BatchCleanupEvent,
    // possibly processed batch job or half-processed batch job need to be processed again,
    // so the batchTime will be equal to lastAllocatedBatchTime.
    // 2. Slow checkpointing makes recovered batch time older than WAL recovered
    // lastAllocatedBatchTime.
    // This situation will only occurs in recovery time.
    logInfo(s"Possibly processed batch $batchTime need to be processed again in WAL recovery")
  }
}

其中，getReceivedBlockQueue的源码如下：

/** Get the queue of received blocks belonging to a particular stream */
private def getReceivedBlockQueue(streamId: Int): ReceivedBlockQueue = {
  streamIdToUnallocatedBlockQueues.getOrElseUpdate(streamId, new ReceivedBlockQueue)
}

可以看到，worker node 发送过来的block 数据被取出来了。

根据WAL block创建 RDD

org.apache.spark.streaming.dstream.ReceiverInputDStream#createBlockRDD 源码如下：

private[streaming] def createBlockRDD(time: Time, blockInfos: Seq[ReceivedBlockInfo]): RDD[T] = {

  if (blockInfos.nonEmpty) {
    val blockIds = blockInfos.map { _.blockId.asInstanceOf[BlockId] }.toArray
   // 所有的block已经有了WriteAheadLogRecordHandle， 创建一个WALBackedBlockRDD即可， 否则创建BlockRDD。
// 其中，WriteAheadLogRecordHandle 是一个跟WAL 相关联的EntryInfo，实现类FileBasedWriteAheadLogSegment就包含了WAL segment 的path， offset 以及 length 信息。RDD 在真正需要数据时，根据这些handle信息从 WAL 中读取数据。
    // Are WAL record handles present with all the blocks
    val areWALRecordHandlesPresent = blockInfos.forall { _.walRecordHandleOption.nonEmpty }

    if (areWALRecordHandlesPresent) {
      // If all the blocks have WAL record handle, then create a WALBackedBlockRDD
      val isBlockIdValid = blockInfos.map { _.isBlockIdValid() }.toArray
      val walRecordHandles = blockInfos.map { _.walRecordHandleOption.get }.toArray
      new WriteAheadLogBackedBlockRDD[T](
        ssc.sparkContext, blockIds, walRecordHandles, isBlockIdValid)
    } else {
      // Else, create a BlockRDD. However, if there are some blocks with WAL info but not
      // others then that is unexpected and log a warning accordingly.
      if (blockInfos.find(_.walRecordHandleOption.nonEmpty).nonEmpty) {
        if (WriteAheadLogUtils.enableReceiverLog(ssc.conf)) {
          logError("Some blocks do not have Write Ahead Log information; " +
            "this is unexpected and data may not be recoverable after driver failures")
        } else {
          logWarning("Some blocks have Write Ahead Log information; this is unexpected")
        }
      }
      val validBlockIds = blockIds.filter { id =>
        ssc.sparkContext.env.blockManager.master.contains(id)
      }
      if (validBlockIds.size != blockIds.size) {
        logWarning("Some blocks could not be recovered as they were not found in memory. " +
          "To prevent such data loss, enabled Write Ahead Log (see programming guide " +
          "for more details.")
      }
      new BlockRDD[T](ssc.sc, validBlockIds)
    }
  } else {
    // If no block is ready now, creating WriteAheadLogBackedBlockRDD or BlockRDD
    // according to the configuration
    if (WriteAheadLogUtils.enableReceiverLog(ssc.conf)) {
      new WriteAheadLogBackedBlockRDD[T](
        ssc.sparkContext, Array.empty, Array.empty, Array.empty)
    } else {
      new BlockRDD[T](ssc.sc, Array.empty)
    }
  }
}

org.apache.spark.streaming.rdd.WriteAheadLogBackedBlockRDD#compute 的源码如下：

/**
 * Gets the partition data by getting the corresponding block from the block manager.
 * If the block does not exist, then the data is read from the corresponding record
 * in write ahead log files.
 */
override def compute(split: Partition, context: TaskContext): Iterator[T] = {
  assertValid()
  val hadoopConf = broadcastedHadoopConf.value
  val blockManager = SparkEnv.get.blockManager
  val partition = split.asInstanceOf[WriteAheadLogBackedBlockRDDPartition]
  val blockId = partition.blockId

  def getBlockFromBlockManager(): Option[Iterator[T]] = {
    blockManager.get(blockId).map(_.data.asInstanceOf[Iterator[T]])
  }

  def getBlockFromWriteAheadLog(): Iterator[T] = {
    var dataRead: ByteBuffer = null
    var writeAheadLog: WriteAheadLog = null
    try {
      // The WriteAheadLogUtils.createLog*** method needs a directory to create a
      // WriteAheadLog object as the default FileBasedWriteAheadLog needs a directory for
      // writing log data. However, the directory is not needed if data needs to be read, hence
      // a dummy path is provided to satisfy the method parameter requirements.
      // FileBasedWriteAheadLog will not create any file or directory at that path.
      // FileBasedWriteAheadLog will not create any file or directory at that path. Also,
      // this dummy directory should not already exist otherwise the WAL will try to recover
      // past events from the directory and throw errors.
      val nonExistentDirectory = new File(
        System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString).getAbsolutePath
      writeAheadLog = WriteAheadLogUtils.createLogForReceiver(
        SparkEnv.get.conf, nonExistentDirectory, hadoopConf)
      dataRead = writeAheadLog.read(partition.walRecordHandle)
    } catch {
      case NonFatal(e) =>
        throw new SparkException(
          s"Could not read data from write ahead log record ${partition.walRecordHandle}", e)
    } finally {
      if (writeAheadLog != null) {
        writeAheadLog.close()
        writeAheadLog = null
      }
    }
    if (dataRead == null) {
      throw new SparkException(
        s"Could not read data from write ahead log record ${partition.walRecordHandle}, " +
          s"read returned null")
    }
    logInfo(s"Read partition data of $this from write ahead log, record handle " +
      partition.walRecordHandle)
    if (storeInBlockManager) {
      blockManager.putBytes(blockId, dataRead, storageLevel)
      logDebug(s"Stored partition data of $this into block manager with level $storageLevel")
      dataRead.rewind()
    }
    blockManager.dataDeserialize(blockId, dataRead).asInstanceOf[Iterator[T]]
  }
 // 如果partition.isBlockIdValid 为true，则说明该 block 数据存在executors 中
  if (partition.isBlockIdValid) {
// 先根据 BlockManager从 executor中读取数据， 如果没有，再从WAL 中读取数据
// BlockManager 从内存还是从磁盘上获取的数据 ？
blockManager 从 local 或 remote 获取 block，其中 local既可以从 memory 中获取也可以从 磁盘中读取， 其中remote获取数据是同步的，即在fetch block 过程中会一直blocking。
    getBlockFromBlockManager().getOrElse { getBlockFromWriteAheadLog() }
  } else {
    getBlockFromWriteAheadLog()
  }
}

至此，从启动 receiver，到receiver 接收数据并保存到WAL block，driver 接收WAL 的block 信息，直到spark streaming 通过WAL RDD 来获取数据等等都一一做了说明。