將 Spark Streaming + Kafka direct 的 offset 保存進入Zookeeper

在上一遍《“Spark Streaming + Kafka direct + checkpoints + 代碼改變” 引發的問題》中說到，當時是將 topic 的 partition 的 offset 保存到了 MySQL 數據庫中，其存在一個問題，就是無法在現有的監控工具中進行體現（如：Kafka Manager）。那我們現在就來將此offset保存到zookeeper中，從而使用監控工具發揮其效果。

一、Kafka Zookeeper 結構

首先，我們來看一下，Kafka在Zookeeper中的結構樹，如圖（說明：圖片來源於網絡）：

此處，我們主要關心的是 /consumers/[groupId]/offsets/[topic]/[partitionId] -> long (offset) 這個路徑，具體含義請看圖片上的說明。
由此可見，我們只需要在原來保存到數據庫中的內容現在寫入到如上路徑中即可。

---

二、具體實現

1、程序實現，如下：

public class SparkStreamingOnKafkaDirect{

    public static JavaStreamingContext createContext(){

        SparkConf conf = new SparkConf().setMaster("local[4]").setAppName("SparkStreamingOnKafkaDirect");

        JavaStreamingContext jsc = new JavaStreamingContext(conf, Durations.seconds(30));
        jsc.checkpoint("/checkpoint");

        Map<String, String> kafkaParams = new HashMap<String, String>();
        kafkaParams.put("metadata.broker.list","192.168.1.151:1234,192.168.1.151:1235,192.168.1.151:1236");

        Set<String> topics = new HashSet<String>();
        topics.add("kafka_direct");

        JavaPairInputDStream<String, String> lines = KafkaUtils.createDirectStream(jsc, String.class,
                        String.class, StringDecoder.class,
                        StringDecoder.class, kafkaParams,
                        topics);

        final AtomicReference<OffsetRange[]> offsetRanges = new AtomicReference<>();

        JavaDStream<String> words = lines.transformToPair(
                new Function<JavaPairRDD<String, String>, JavaPairRDD<String, String>>() {
                    @Override
                    public JavaPairRDD<String, String> call(JavaPairRDD<String, String> rdd) throws Exception {
                      OffsetRange[] offsets = ((HasOffsetRanges) rdd.rdd()).offsetRanges();
                      offsetRanges.set(offsets);
                      return rdd;
                    }
                  }
                ).flatMap(new FlatMapFunction<Tuple2<String, String>, String>() {
                    public Iterable<String> call(
                            Tuple2<String, String> event)
                            throws Exception {
                        String line = event._2;
                        return Arrays.asList(line);
                    }
                });

        JavaPairDStream<String, Integer> pairs = words
                .mapToPair(new PairFunction<String, String, Integer>() {

                    public Tuple2<String, Integer> call(
                            String word) throws Exception {
                        return new Tuple2<String, Integer>(
                                word, 1);
                    }
                });

        JavaPairDStream<String, Integer> wordsCount = pairs
                .reduceByKey(new Function2<Integer, Integer, Integer>() {
                    public Integer call(Integer v1, Integer v2)
                            throws Exception {
                        return v1 + v2;
                    }
                });

        lines.foreachRDD(new VoidFunction<JavaPairRDD<String,String>>(){
            @Override
            public void call(JavaPairRDD<String, String> t) throws Exception {

                ObjectMapper objectMapper = new ObjectMapper();

                CuratorFramework  curatorFramework = CuratorFrameworkFactory.builder()
                        .connectString("192.168.1.151:2181").connectionTimeoutMs(1000)
                        .sessionTimeoutMs(10000).retryPolicy(new RetryUntilElapsed(1000, 1000)).build();

                curatorFramework.start();

                for (OffsetRange offsetRange : offsetRanges.get()) {
                    final byte[] offsetBytes = objectMapper.writeValueAsBytes(offsetRange.untilOffset());
                    String nodePath = "/consumers/spark-group/offsets/" + offsetRange.topic()+ "/" + offsetRange.partition();
                    if(curatorFramework.checkExists().forPath(nodePath)!=null){
                            curatorFramework.setData().forPath(nodePath,offsetBytes);
                        }else{
                            curatorFramework.create().creatingParentsIfNeeded().forPath(nodePath, offsetBytes);
                        }
                }
            }
            curatorFramework.close();
        });

        wordsCount.print();

        return jsc;
    }

    public static void main(String[] args) {
        JavaStreamingContextFactory factory = new JavaStreamingContextFactory() {
            public JavaStreamingContext create() {
              return createContext();
            }
          };

        JavaStreamingContext jsc = JavaStreamingContext.getOrCreate("/checkpoint", factory);

        jsc.start();

        jsc.awaitTermination();
        jsc.close();
    }

}

2、準備測試環境，並記錄目前consumer中的信息，如下圖：

從截圖中可以看出，目前 kafka_topic 這個 topic 中，所有的消息都已經被消息了。
說明：由於在寫blog之前，我已經運行過代碼了，所以可以直接看到數據，如果是第一次運行，則是看到不 spark-group 這個 group 的。

3、運行Spark Streaming 程序(注意：要先清空 checkpoint 目錄下的內容)，在程序運行的過程中查看 kafka monitor中關於 spark-group的變化情況：
第 n 次 Job 之後：

從圖片上可以看出，已經起到作用了。