spark-streaming-kafka-0-10源码分析

转发请注明原创地址http://www.cnblogs.com/dongxiao-yang/p/7767621.html

本文所研究的spark-streaming代码版本为2.3.0-SNAPSHOT

spark-streaming为了匹配0.10以后版本的kafka客户端变化推出了一个目前还是Experimental状态的spark-streaming-kafka-0-10客户端，由于老的0.8版本无法支持kerberos权限校验，需要研究下spark-streaming-kafka-0-10的源码实现以及系统架构。

首先看下初始化kafkastream的方法声明，

def createDirectStream[K, V](
ssc: StreamingContext,
locationStrategy: LocationStrategy,
consumerStrategy: ConsumerStrategy[K, V],
perPartitionConfig: PerPartitionConfig
): InputDStream[ConsumerRecord[K, V]] = {
new DirectKafkaInputDStream[K, V](ssc, locationStrategy, consumerStrategy, perPartitionConfig)
}

DirectKafkaInputDStream的初始化参数包括StreamingContext，LocationStrategy，ConsumerStrategy和perPartitionConfig，根据源码文档locationStrategy一般采用PreferConsistent

，perPartitionConfig一般采用默认实现，这里不做研究，主要会有点区别的参数为consumerStrategy，它的作用会在下面的源码分析里展示出来。

一  driver consumer

JavaInputDStream<ConsumerRecord<String, String>> stream = KafkaUtils
                .createDirectStream(jssc,
                        LocationStrategies.PreferConsistent(),
                        ConsumerStrategies.<String, String> Subscribe(topics,
                                kafkaParams));

以上述初始化代码为例，首先DirectKafkaInputDStream会调用start方法进行初始化，相关代码如下

  override def start(): Unit = {
    val c = consumer //初始化driver端consumer
    paranoidPoll(c)  //调整offset位置
    if (currentOffsets.isEmpty) {
      currentOffsets = c.assignment().asScala.map { tp =>
        tp -> c.position(tp)
      }.toMap
    }

    // don't actually want to consume any messages, so pause all partitions
    c.pause(currentOffsets.keySet.asJava)
  }

这段代码在driver端初始化一个consumer, 该consumer的类型由上面提到的consumerStrategy决定，Subscribe类的实现如下，相当与在driver端启动一个以subscribe模式订阅topic的客户端。在有初始启动offset传入的情况下会把consumer的offset游标seek到对应的地址。

private case class Subscribe[K, V](
    topics: ju.Collection[jl.String],
    kafkaParams: ju.Map[String, Object],
    offsets: ju.Map[TopicPartition, jl.Long]
  ) extends ConsumerStrategy[K, V] with Logging {

  def executorKafkaParams: ju.Map[String, Object] = kafkaParams

  def onStart(currentOffsets: ju.Map[TopicPartition, jl.Long]): Consumer[K, V] = {
    val consumer = new KafkaConsumer[K, V](kafkaParams)
    consumer.subscribe(topics)
    val toSeek = if (currentOffsets.isEmpty) {
      offsets
    } else {
      currentOffsets
    }
    if (!toSeek.isEmpty) {
      // work around KAFKA-3370 when reset is none
      // poll will throw if no position, i.e. auto offset reset none and no explicit position
      // but cant seek to a position before poll, because poll is what gets subscription partitions
      // So, poll, suppress the first exception, then seek
      val aor = kafkaParams.get(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG)
      val shouldSuppress =
        aor != null && aor.asInstanceOf[String].toUpperCase(Locale.ROOT) == "NONE"
      try {
        consumer.poll(0)
      } catch {
        case x: NoOffsetForPartitionException if shouldSuppress =>
          logWarning("Catching NoOffsetForPartitionException since " +
            ConsumerConfig.AUTO_OFFSET_RESET_CONFIG + " is none.  See KAFKA-3370")
      }
      toSeek.asScala.foreach { case (topicPartition, offset) =>
          consumer.seek(topicPartition, offset)
      }
      // we've called poll, we must pause or next poll may consume messages and set position
      consumer.pause(consumer.assignment())
    }

    consumer
  }
}

DirectKafkaInputDStream的另一个核心方法是compute，这个方法的核心作用之一就是不断地生成对应时间的RDD分配到新的job计算任务，具体实现如下，主要是根据系统设置的限速和现有

kafka topicpartion计算出每一个job分配到的KafkaRDD对应的数据范围以及提交offset等工作。

  override def compute(validTime: Time): Option[KafkaRDD[K, V]] = {
    val untilOffsets = clamp(latestOffsets()) //根据maxrate和backpressuce等限速配置计算下一批rdd每个里面kafka消息的截止offset
    val offsetRanges = untilOffsets.map { case (tp, uo) =>
      val fo = currentOffsets(tp)
      OffsetRange(tp.topic, tp.partition, fo, uo)
    }//初始化offset列表，包括(topic,partition,起始offset，截止offset)
    val useConsumerCache = context.conf.getBoolean("spark.streaming.kafka.consumer.cache.enabled",
      true)
    val rdd = new KafkaRDD[K, V](context.sparkContext, executorKafkaParams, offsetRanges.toArray,
      getPreferredHosts, useConsumerCache)//根据计算好的offsetRange和修改后的kafkaParam初始化RDD

    // Report the record number and metadata of this batch interval to InputInfoTracker.
    val description = offsetRanges.filter { offsetRange =>
      // Don't display empty ranges.
      offsetRange.fromOffset != offsetRange.untilOffset
    }.map { offsetRange =>
      s"topic: ${offsetRange.topic}\tpartition: ${offsetRange.partition}\t" +
        s"offsets: ${offsetRange.fromOffset} to ${offsetRange.untilOffset}"
    }.mkString("\n")
    // Copy offsetRanges to immutable.List to prevent from being modified by the user
    val metadata = Map(
      "offsets" -> offsetRanges.toList,
      StreamInputInfo.METADATA_KEY_DESCRIPTION -> description)
    val inputInfo = StreamInputInfo(id, rdd.count, metadata)
    ssc.scheduler.inputInfoTracker.reportInfo(validTime, inputInfo)

    currentOffsets = untilOffsets
    commitAll()
    Some(rdd)
  }

注意上文里的latestOffset()方法实现如下，通过新的consumerapi的c.seekToEnd(currentOffsets.keySet.asJava)将consumer的offsetapi游标放到了对应分区的最后位置，

如果在初始化的kafkaParams设置"enable.auto.commit"属性为"true"，diver客户端会自动像kafka发送最后seek到的offset位置

  protected def latestOffsets(): Map[TopicPartition, Long] = {
    val c = consumer
    paranoidPoll(c)
    val parts = c.assignment().asScala

    // make sure new partitions are reflected in currentOffsets
    val newPartitions = parts.diff(currentOffsets.keySet)
    // position for new partitions determined by auto.offset.reset if no commit
    currentOffsets = currentOffsets ++ newPartitions.map(tp => tp -> c.position(tp)).toMap
    // don't want to consume messages, so pause
    c.pause(newPartitions.asJava)
    // find latest available offsets
    c.seekToEnd(currentOffsets.keySet.asJava)
    parts.map(tp => tp -> c.position(tp)).toMap
  }

二  executor consumer

executor consumer的初始化过程位于KafkaRDD内部，在程序初始的kafaparams基础上调用了fixKfkaParams方法对参数进行了部分调整和改写，包括groupid，enable.auto.commit,auto.offset.config等属性。

  private[kafka010] def fixKafkaParams(kafkaParams: ju.HashMap[String, Object]): Unit = {
    logWarning(s"overriding ${ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG} to false for executor")
    kafkaParams.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false: java.lang.Boolean)

    logWarning(s"overriding ${ConsumerConfig.AUTO_OFFSET_RESET_CONFIG} to none for executor")
    kafkaParams.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "none")

    // driver and executor should be in different consumer groups
    val originalGroupId = kafkaParams.get(ConsumerConfig.GROUP_ID_CONFIG)
    if (null == originalGroupId) {
      logError(s"${ConsumerConfig.GROUP_ID_CONFIG} is null, you should probably set it")
    }
    val groupId = "spark-executor-" + originalGroupId
    logWarning(s"overriding executor ${ConsumerConfig.GROUP_ID_CONFIG} to ${groupId}")
    kafkaParams.put(ConsumerConfig.GROUP_ID_CONFIG, groupId)

    // possible workaround for KAFKA-3135
    val rbb = kafkaParams.get(ConsumerConfig.RECEIVE_BUFFER_CONFIG)
    if (null == rbb || rbb.asInstanceOf[java.lang.Integer] < 65536) {
      logWarning(s"overriding ${ConsumerConfig.RECEIVE_BUFFER_CONFIG} to 65536 see KAFKA-3135")
      kafkaParams.put(ConsumerConfig.RECEIVE_BUFFER_CONFIG, 65536: java.lang.Integer)
    }
  }

KafkaRDD内部与consumer相关的几个方法如下：首先通过getPartitions方法将对应的topic分区与RDD的每一个分区对应起来，然后通过compute方法初始化KafkaRDDIterator，每个KafkaRDDIterator通过CachedKafkaConsumer接口拿到一个CachedKafkaConsumer引用并在next()方法里不断返回ConsumerRecord值。

  override def getPartitions: Array[Partition] = {
    offsetRanges.zipWithIndex.map { case (o, i) =>
        new KafkaRDDPartition(i, o.topic, o.partition, o.fromOffset, o.untilOffset)
    }.toArray
  }

override def compute(thePart: Partition, context: TaskContext): Iterator[ConsumerRecord[K, V]] = {
    val part = thePart.asInstanceOf[KafkaRDDPartition]
    assert(part.fromOffset <= part.untilOffset, errBeginAfterEnd(part))
    if (part.fromOffset == part.untilOffset) {
      logInfo(s"Beginning offset ${part.fromOffset} is the same as ending offset " +
        s"skipping ${part.topic} ${part.partition}")
      Iterator.empty
    } else {
      new KafkaRDDIterator(part, context)
    }
  }

  private class KafkaRDDIterator(
      part: KafkaRDDPartition,
      context: TaskContext) extends Iterator[ConsumerRecord[K, V]] {

    logInfo(s"Computing topic ${part.topic}, partition ${part.partition} " +
      s"offsets ${part.fromOffset} -> ${part.untilOffset}")

    val groupId = kafkaParams.get(ConsumerConfig.GROUP_ID_CONFIG).asInstanceOf[String]

    context.addTaskCompletionListener{ context => closeIfNeeded() }

    val consumer = if (useConsumerCache) {
      CachedKafkaConsumer.init(cacheInitialCapacity, cacheMaxCapacity, cacheLoadFactor)
      if (context.attemptNumber >= 1) {
        // just in case the prior attempt failures were cache related
        CachedKafkaConsumer.remove(groupId, part.topic, part.partition)
      }
      CachedKafkaConsumer.get[K, V](groupId, part.topic, part.partition, kafkaParams)
    } else {
      CachedKafkaConsumer.getUncached[K, V](groupId, part.topic, part.partition, kafkaParams)
    }

    var requestOffset = part.fromOffset

    def closeIfNeeded(): Unit = {
      if (!useConsumerCache && consumer != null) {
        consumer.close
      }
    }

    override def hasNext(): Boolean = requestOffset < part.untilOffset

    override def next(): ConsumerRecord[K, V] = {
      assert(hasNext(), "Can't call getNext() once untilOffset has been reached")
      val r = consumer.get(requestOffset, pollTimeout)
      requestOffset += 1
      r
    }
  }

根据是否使用consumer的缓存池特性（这个属性由spark.streaming.kafka.consumer.cache.enabled决定），CachedKafkaConsumer提供了两种静态方法获取consumer客户端，get()和getUncached()。

get方法从CachedKafkaConsumer的静态linkhashmap属性cache中存取已经初始化好的CachedKafkaConsumer对象，相当于每个executor内部维护了一个consumer的连接池。

getUncached相当于每次拉新数据都初始化一个consumer连接，并在这个RDD任务结束后关掉consumer实例。

CachedKafkaConsumer初始化kafka consumer客户端的相关代码如下，可以看到真正拉数据的executor客户端是采用了assgin方式订阅到单个分区初始化完成的。

  protected val consumer = {
    val c = new KafkaConsumer[K, V](kafkaParams)
    val tps = new ju.ArrayList[TopicPartition]()
    tps.add(topicPartition)
    c.assign(tps)
    c
  }

三 offset提交

除了上文提到的将driver端的auto.commit属性打开提交offset的方式以外，sparkstreaming还在DirectKafkaInputDStream中提供了一个commitAsync(offsetRanges: Array[OffsetRange], callback: OffsetCommitCallback)方法允许手动触发offset提交，这个方法将需要提交的offset列表放到了一个commitQueue里面，然后在每次调用compute方法的时候最后的commitall方法通过driver端的consumer把offset提交到kafka上。

  def commitAsync(offsetRanges: Array[OffsetRange], callback: OffsetCommitCallback): Unit = {
    commitCallback.set(callback)
    commitQueue.addAll(ju.Arrays.asList(offsetRanges: _*))
  }

  protected def commitAll(): Unit = {
    val m = new ju.HashMap[TopicPartition, OffsetAndMetadata]()
    var osr = commitQueue.poll()
    while (null != osr) {
      val tp = osr.topicPartition
      val x = m.get(tp)
      val offset = if (null == x) { osr.untilOffset } else { Math.max(x.offset, osr.untilOffset) }
      m.put(tp, new OffsetAndMetadata(offset))
      osr = commitQueue.poll()
    }
    if (!m.isEmpty) {
      consumer.commitAsync(m, commitCallback.get)
    }
  }

stream.foreachRDD(rdd -> {
  OffsetRange[] offsetRanges = ((HasOffsetRanges) rdd.rdd()).offsetRanges();

  // some time later, after outputs have completed
  ((CanCommitOffsets) stream.inputDStream()).commitAsync(offsetRanges);
});

注意：如果是采用官方文档里上述方式手动提交offset，需要把stream对象的属性标记为static或者transient避免序列化，不然可能在任务提交的时候报DirectKafkaInputDStream 无法序列化导致Task not serializable错误

结论

新的spark-streaming-kafka-0-10客户端采用了与原有版本完全不同的架构，一个job里面运行了两组consumer：driver consumer和 executor consumer，driver端consumer负责分配和提交offset到初始化好的KafkaRDD当中去，KafkaRDD内部会根据分配到的每个topic的每个partition初始化一个CachedKafkaConsumer客户端通过assgin的方式订阅到topic拉取数据。