SparkStreaming自适应上游kafka topic partition数目变化

背景
Spark Streaming 作业在运行过程中，上游 topic 增加 partition 数目从 A 增加到 B，会造成作业丢失数据，因为该作业只从 topic 中读取了原来的 A 个 partition 的数据，新增的 B-A 个 partition 的数据会被忽略掉。


思考过程
为了作业能够长时间的运行，一开始遇到这种情况的时候，想到两种方案：


感知上游 topic 的 partition 数目变化，然后发送报警，让用户重启
直接在作业内部自适应上游 topic partition 的变化，完全不影响作业
方案 1 是简单直接，第一反应的结果，但是效果不好，需要用户人工介入，而且需要删除 checkpoint 文件


方案 2 从根本上解决问题，用户不需要关心上游 partition 数目的变化，但是第一眼会觉得较难实现。


方案 1 很快被 pass 掉，因为人工介入的成本太高，而且实现起来很别扭。接下来考虑方案 2.


Spark Streaming 程序中使用 Kafka 的最原始方式为 KafkaUtils.createDirectStream 通过源码，我们找到调用链条大致是这样的


KafkaUtils.createDirectStream   -->   new DirectKafkaInputDStream --> 最终由 DirectKafkaInputDStream#compute(validTime : Time) 函数来生成 KafkaRDD。


而 KafkaRDD 的 partition 数和 作业开始运行时 topic 的 partition 数一致，topic 的 partition 数保存在 currentOffsets 变量中，currentOffsets 是一个 Map[TopicAndPartition, Long]类型的变量，保存每个 partition 当前消费的 offset 值，但是作业运行过程中 currentOffsets 不会增加 key，就是说不会增加 KafkaRDD 的 partition，这样导致每次生成 KafkaRDD 的时候都使用 作业开始运行时 topic 的 partition 数作为 KafkaRDD 的 partition 数，从而会造成数据的丢失。


解决方案
我们只需要在每次生成 KafkaRDD 的时候，将 currentOffsets 修正为正常的值（往里面增加对应的 partition 数，总共 B-A 个，以及每个增加的 partition 的当前 offset 从零开始）。


第一个问题出现了，我们不能修改 Spark 的源代码，重新进行编译，因为这不是我们自己维护的。想到的一种方案是继承 DirectKafkaInputDStream。我们发现不能继承 DirectKafkaInputDStream 该类，因为这个类是使用 private[streaming] 修饰的。
第二个问题出现了，怎么才能够继承 DirectKafkaInputDStream，这时我们只需要将希望继承 DirectKafkaInputDStream 的类放到一个单独的文件 F 中，文件 F 使用 package org.apache.spark.streaming 进行修饰即可，这样可以绕过不能继承 DirectKafkaInputDStream 的问题。这个问题解决后，我们还需要修改 Object KafkaUtils，让该 Object 内部调用我们修改后的 DirectKafkaInputDStream（我命名为 MTDirectKafkaInputDStream)
第三个问题如何让 Spark 调用 MTDirectKafkaInputDStream，而不是 DirectKafkaInputDStream，这里我们使用简单粗暴的方式，将 KafkaUtils 的代码 copy 一份，然后将其中调用 DirectKafkaInputDStream 的部分都修改为 MTDirectKafkaInputDStream，这样就实现了我们的需要。当然该文件也需要使用 package org.apache.spark.streaming 进行修饰
第二个和第三个问题的解决方案在  中国 Spark 技术峰会 2016  上，广点通的 林立伟 有提及，后续会进行尝试
总结下，我们需要做两件事


修改 DirectKafkaInputDStream#compute 使得能够自适应 topic 的 partition 变更
修改 KafkaUtils，使得我们能够调用修改过后的 DirectKafkaInputDStream
代码

package org.apache.spark.streaming.kafka.mt

import com.meituan.data.util.Constants
import com.meituan.service.inf.kms.client.Kms
import kafka.common.{ErrorMapping, TopicAndPartition}
import kafka.javaapi.{TopicMetadata, TopicMetadataRequest}
import kafka.javaapi.consumer.SimpleConsumer
import kafka.message.MessageAndMetadata
import kafka.serializer.Decoder
import org.apache.spark.streaming.{StreamingContext, Time}
import org.apache.spark.streaming.kafka.{DirectKafkaInputDStream, KafkaRDD}

import scala.collection.JavaConverters._
import scala.util.control.Breaks._
import scala.reflect.ClassTag

/**
  * Created by qiucongxian on 10/27/16.
  */
class MTDirectKafkaInputDStream[
  K: ClassTag,
  V: ClassTag,
  U <: Decoder[K]: ClassTag,
  T <: Decoder[V]: ClassTag,
  R: ClassTag](
    @transient ssc_ : StreamingContext,
    val MTkafkaParams: Map[String, String],
    val MTfromOffsets: Map[TopicAndPartition, Long],
    messageHandler: MessageAndMetadata[K, V] => R
) extends DirectKafkaInputDStream[K, V, U, T, R](ssc_, MTkafkaParams , MTfromOffsets, messageHandler) {
    private val kafkaBrokerList : String = "host1:port1,host2:port2,host3:port3" //根据自己的情况自行修改

    override def compute(validTime: Time) : Option[KafkaRDD[K, V, U, T, R]] = {
      /**
        * 在这更新 currentOffsets 从而做到自适应上游 partition 数目变化
        */
        updateCurrentOffsetForKafkaPartitionChange()
        super.compute(validTime)
    }

    private def updateCurrentOffsetForKafkaPartitionChange() : Unit = {
      val topic = currentOffsets.head._1.topic
      val nextPartitions : Int = getTopicMeta(topic) match {
          case Some(x) => x.partitionsMetadata.size()
          case _ => 0
      }
      val currPartitions = currentOffsets.keySet.size

      if (nextPartitions > currPartitions) {
        var i = currPartitions
        while (i < nextPartitions) {
           currentOffsets = currentOffsets + (TopicAndPartition(topic, i) -> 0)
           i = i + 1
        }
      }
      logInfo(s"######### ${nextPartitions}  currentParttions ${currentOffsets.keySet.size} ########")
    }

    private def getTopicMeta(topic: String) : Option[TopicMetadata] = {
        var metaData : Option[TopicMetadata] = None
        var consumer : Option[SimpleConsumer] = None

        val topics = List[String](topic)
        val brokerList = kafkaBrokerList.split(",")
        brokerList.foreach(
          item => {
            val hostPort = item.split(":")
            try {
              breakable {
                  for (i <- 0 to 3) {
                      consumer = Some(new SimpleConsumer(host = hostPort(0), port = hostPort(1).toInt,
                                            soTimeout = 10000, bufferSize = 64 * 1024, clientId = "leaderLookup"))
                      val req : TopicMetadataRequest = new TopicMetadataRequest(topics.asJava)
                      val resp = consumer.get.send(req)

                      metaData = Some(resp.topicsMetadata.get(0))
                      if (metaData.get.errorCode == ErrorMapping.NoError) break()
                  }
              }
            } catch {
              case e => logInfo(s" ###### Error in MTDirectKafkaInputDStream ${e} ######")
            }
          }
        )
        metaData
    }
}

在修改过后的 KafkaUtils 文件中，将所有的 DirectKafkaInputDStream 都替换为 MTDirectKafkaInputDStream 即可。