Spark源码分析——DirectInputDStream
在Spark-Streaming中，对流的抽象是使用DStream来定义的，想要理解Spark-Streaming的流处理模型，深入了解DStream是很有必要的。
DStream
我们在定义一个流的处理逻辑时，首先从一个数据的流入源开始，这个数据源使用InputDStream定义，它是DStream的一个子类，之后我们会在其上调用一些tranform类型算子，像map,reduce,filter等等，每调用一个算子，都会创建一个新的DStream，每一个新创建的DStream都保留着当前节点所依赖的上一个节点，当前节点的执行逻辑这两个部分，这样，多个DStream节点就构成了一个逻辑执行链，这是典型的职责链设计模式。而当调用action类型的算子时，最后底层都会回归到ForEachDStream上
private[streaming]
class ForEachDStream[T: ClassTag] (
parent: DStream[T],
foreachFunc: (RDD[T], Time) => Unit,
displayInnerRDDOps: Boolean
) extends DStream[Unit](parent.ssc) {
override def dependencies: List[DStream[_]] = List(parent)
override def slideDuration: Duration = parent.slideDuration
override def compute(validTime: Time): Option[RDD[Unit]] = None
override def generateJob(time: Time): Option[Job] = {
parent.getOrCompute(time) match {
case Some(rdd) =>
val jobFunc = () => createRDDWithLocalProperties(time, displayInnerRDDOps) {
foreachFunc(rdd, time)
}
Some(new Job(time, jobFunc))
case None => None
}
}
}
这里我们关注generateJob方法，这里调用了它依赖的父DStream的getOrCompute来生成一个需要它来处理的RDD，然后对该RDD做该节点本身需要做的一些操作，即foreachFunc闭包，其实所有DStream的getOrCompute方法底层都会调用compute方法，即所有的DStream的compute方法中，要么其本身能够从外部拉取数据，即InputDStream作为DStream链的第一个节点，要么其本身调用依赖的上游DStream的compute方法，再对生成的RDD做其本节点所定义的一些操作作为其返回值。如此，当DStream链的最后一个节点被调用了compute方法时，它能够依次递归的逐节点的compute方法，最后调用第一个InputDStream节点的compute方法生成一个能够拉取外部数据的RDD。
以上只是为了直观的理解DStream链是如何工作的，具体体现在分布式环境上时，是由RDD来定义操作，切分成task后由Executor来执行。
另外需要说的是如果我们在单个流上定义一系列除window外的操作，其和我们直接调用InputDStream的foreachRDD后，在rdd上定义操作是等效的。
总体模型
除了上面介绍的DStream之外，在Spark-Streaming内部还有一些保存作业处理逻辑的模块和用于根据时间生成和管理每个批次数据的模块。下面是在SparkStreaming中一些比较核心的类，他们是构成一个流式作业，和使其运行起来的框架。
- InputDStream 管理流的数据源的通用抽象类
- JobGenerator 作业生成器
- JobScheduler 作业调度器，用于提交作业到集群运行
- DStreamGraph 管理创建的所有InputDStream的初始化和启动，但不负责InputDStream间依赖关系的管理，InputDStream间依赖关系由其子类实现管理
在我们不了解以上类如何交互前，是没办法读懂DirectInputDStream的源码的，这里我们简要介绍一下这四个类之间是如何交互的，想要了解详细流程的请阅读我的另一篇博客(嘿嘿，这广告打的不要不要的~~)
JobGenerator每隔我们设定的时间间隔会生成一个事件用于触发生成一个作业，生成作业的过程是DStreamGraph实现的，他会返回一个Job对象，里面包含了一个需要在其上执行计算的RDD，包括所有计算逻辑的闭包，而这个闭包真正执行，是在JobScheduler将这个Job对象提交到一个线程池之后，其会在线程池内执行这个Job对象内的闭包逻辑，将其转换成分布式计算的task分发到不同的节点上去执行。
我们从图中看到其底层是通过执行DStream的compute方法创建的RDD，而compute的具体执行逻辑是由DStream的子类负责实现决定的。我们把这里作为我们阅读DirectInputDStream的入口。但是在阅读这段代码之前，我们先看一下DirectInputDStream是怎么被创建的，DirectInputDStream的构造函数是私有的，只能使用KafkaUtils来创建，代码如下：
def createDirectStream[K, V](
ssc: StreamingContext,
locationStrategy: LocationStrategy,
consumerStrategy: ConsumerStrategy[K, V]
): InputDStream[ConsumerRecord[K, V]] = {
new DirectKafkaInputDStream[K, V](ssc, locationStrategy, consumerStrategy, perPartitionConfig)
}
第一个参数是流上下文，在构造InputDStream时需要传递进去，第二个参数是LocationStrategy，用于在调度task的时候提供一个优化选项，第三个参数是ConsumerStrategy，用于指定消费模式，是自动订阅，还是手动指定消费的partition，同时可以指定消费的起始offset，和一些额外的kafka参数。
DirectInputDStream会使用ConsumerStrategy中的一些参数在driver端创建一个consumer，用来获取每个批次需要消费到的最新的offset和提交已经消费过的消息的offset。
val executorKafkaParams = {
val ekp = new ju.HashMap[String, Object](consumerStrategy.executorKafkaParams)
KafkaUtils.fixKafkaParams(ekp)
ekp
}
protected var currentOffsets = Map[TopicPartition, Long]()
@transient private var kc: Consumer[K, V] = null
def consumer(): Consumer[K, V] = this.synchronized {
if (null == kc) {
kc = consumerStrategy.onStart(currentOffsets.mapValues(l => new java.lang.Long(l)).asJava)
}
kc
}
而ConsumerStrategy的onStart方法我们选择一个比较常用的Subscribe类来看：
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) {
consumer.poll(0)
toSeek.asScala.foreach { case (topicPartition, offset) =>
consumer.seek(topicPartition, offset)
}
consumer.pause(consumer.assignment())
}
consumer
}
我们看到在onStart方法中创建了一个kafka consumer，并从传递进来的方法参数和构造参数中优先选择方法参数来对起始offset进行重置。因为当DirectInputDStream第一次调动onStart方法时候其传递进来的map就是一个空的map，此时确实应该使用Subscribe类本身的起始offset，而如果Subscribe的offset也为空时，则启用kafka集群内部存储的消费进度。
好了，了解了DirectInputDStream在被创建时做的一些初始化后，我们回过头来看compute的代码，代码如下：
override def compute(validTime: Time): Option[KafkaRDD[K, V]] = {
val untilOffsets = clamp(latestOffsets())
val offsetRanges = untilOffsets.map { case (tp, uo) =>
val fo = currentOffsets(tp)
OffsetRange(tp.topic, tp.partition, fo, uo)
}
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)
currentOffsets = untilOffsets
commitAll()
Some(rdd)
}
这里我们省去一些不必要的代码，只保留核心逻辑代码。基本的逻辑就是，获取本次应该消费到的offset，确定每个分区应该处理的数据范围，创建KafkaRDD，提交所有之前的offset，返回RDD，我们关注一下clamp和latestOffsets方法。代码如下：
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
}
protected def clamp(
offsets: Map[TopicPartition, Long]): Map[TopicPartition, Long] = {
maxMessagesPerPartition(offsets).map { mmp =>
mmp.map { case (tp, messages) =>
val uo = offsets(tp)
tp -> Math.min(currentOffsets(tp) + messages, uo)
}
}.getOrElse(offsets)
}
在latestOffsets中得到最新的需要消费到的offset，同时如果有新的分区，将新的分区加入进来。 clamp用于修正针对有限速情况下的offset，限速方式通过spark.streaming.kafka.maxRatePerPartition配置项来指定
然后我们回过头来看创建KafkaRDD的部分
new KafkaRDD[K, V](context.sparkContext, executorKafkaParams, offsetRanges.toArray, getPreferredHosts, useConsumerCache)
来看看其构造函数的参数，SparkContext，在创建父类RDD的时候需要该参数，executorKafkaParams是在Executor端创建KafkaConsumer进行消费的时候需要的一些配置参数，由用户通过KafkaUtils创建DirectInputDStream传递进来，offsetRanges.toArray，指明了该RDD内包含的每个分区的消息消费范围，getPreferredHosts，这个是用于优化决定在提交计算时，每个Task应该被分配到哪台机器上而提供的优化参数。最后一个就是useConsumerCache参数，用于指示Executor端在拉取数据的时候是否会缓存和复用KafkaConsumer。
下面简单介绍一下KafkaRDD，我们主要看两个方法：
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]
require(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 {
logInfo(s"Computing topic ${part.topic}, partition ${part.partition} " +
s"offsets ${part.fromOffset} -> ${part.untilOffset}")
if (compacted) {
new CompactedKafkaRDDIterator[K, V](
part,
context,
kafkaParams,
useConsumerCache,
pollTimeout,
cacheInitialCapacity,
cacheMaxCapacity,
cacheLoadFactor
)
} else {
new KafkaRDDIterator[K, V](
part,
context,
kafkaParams,
useConsumerCache,
pollTimeout,
cacheInitialCapacity,
cacheMaxCapacity,
cacheLoadFactor
)
}
}
}
首先getPartitions方法将要消费的分区转化为KafkaRDDPartition，这部分没啥说的，这个getPartitions是在RDD类的partitions方法中被调用到的。
另一个compute方法是RDD中最重要的方法，该方法根据分区和TaskContext返回一个该分区数据的迭代器，这里当分区的起始和终止offset不相同的时候，会根据spark.streaming.kafka.allowNonConsecutiveOffsets这个的配置项来决定是返回CompactedKafkaRDDIterator还是KafkaRDDIterator，这个选项是在Spark 2.3版本中后加进去的，为了解决之前由于在Kafka中配置的某个topic数据是compact的而导致的offset不连续问题。
通过比对二者的next方法就可以看出来： KafkaRDDIterator.next:
override def next(): ConsumerRecord[K, V] = {
if (!hasNext) {
throw new ju.NoSuchElementException("Can't call getNext() once untilOffset has been reached")
}
val r = consumer.get(requestOffset, pollTimeout)
requestOffset += 1
r
}
CompactedKafkaRDDIterator.next:
override def next(): ConsumerRecord[K, V] = {
if (!hasNext) {
throw new ju.NoSuchElementException("Can't call getNext() once untilOffset has been reached")
}
val r = nextRecord
if (r.offset + 1 >= part.untilOffset) {
okNext = false
} else {
nextRecord = consumer.compactedNext(pollTimeout)
if (nextRecord.offset >= part.untilOffset) {
okNext = false
consumer.compactedPrevious()
}
}
r
}
//compactedNext方法
def compactedNext(pollTimeoutMs: Long): ConsumerRecord[K, V] = {
if (!buffer.hasNext()) {
poll(pollTimeoutMs)
}
require(buffer.hasNext(),
s"Failed to get records for compacted $groupId $topicPartition " +
s"after polling for $pollTimeoutMs")
val record = buffer.next()
nextOffset = record.offset + 1
record
}
可以明显看出，KafkaRDDIterator的next方法是要求offset是严格按序递增的，而CompactedKafkaRDDIterator的next方法则没有这方面的限制，只需要依次获取下一个record就可以了。
好了，以上就是对DirectInputDStream及其相关代码的源码分析，如有错误，欢迎指出。
1万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
