<commit_before>/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package kafka.tools


import joptsimple._
import org.I0Itec.zkclient.ZkClient
import kafka.utils.{ZkUtils, ZKStringSerializer, Logging}
import kafka.consumer.SimpleConsumer
import kafka.api.{PartitionOffsetRequestInfo, OffsetRequest}
import kafka.common.TopicAndPartition
import scala.collection._


object ConsumerOffsetChecker extends Logging {

  private val consumerMap: mutable.Map[String, Option[SimpleConsumer]] = mutable.Map()

  private val BidPidPattern = """(\d+)-(\d+)""".r

  private val BrokerIpPattern = """.*:([^:]+):(\d+$)""".r
  // e.g., 127.0.0.1-1315436360737:127.0.0.1:9092
  // e.g., host.domain.com-1315436360737:host.domain.com:9092

  private def getConsumer(zkClient: ZkClient, bid: String): Option[SimpleConsumer] = {
    val brokerInfo = ZkUtils.readDataMaybeNull(zkClient, "/brokers/ids/%s".format(bid))._1
    val consumer = brokerInfo match {
      case BrokerIpPattern(ip, port) =>
        Some(new SimpleConsumer(ip, port.toInt, 10000, 100000, "ConsumerOffsetChecker"))
      case _ =>
        error("Could not parse broker info %s".format(brokerInfo))
        None
    }
    consumer
  }

  private def processPartition(zkClient: ZkClient,
                               group: String, topic: String, bidPid: String) {
    val offset = ZkUtils.readData(zkClient, "/consumers/%s/offsets/%s/%s".
            format(group, topic, bidPid))._1.toLong
    val owner = ZkUtils.readDataMaybeNull(zkClient, "/consumers/%s/owners/%s/%s".
            format(group, topic, bidPid))._1
    println("%s,%s,%s (Group,Topic,BrokerId-PartitionId)".format(group, topic, bidPid))
    println("%20s%s".format("Owner = ", owner))
    println("%20s%d".format("Consumer offset = ", offset))
    println("%20s%,d (%,.2fG)".format("= ", offset, offset / math.pow(1024, 3)))

    bidPid match {
      case BidPidPattern(bid, pid) =>
        val consumerOpt = consumerMap.getOrElseUpdate(
          bid, getConsumer(zkClient, bid))
        consumerOpt match {
          case Some(consumer) =>
            val topicAndPartition = TopicAndPartition(topic, pid.toInt)
            val request =
              OffsetRequest(immutable.Map(topicAndPartition -> PartitionOffsetRequestInfo(OffsetRequest.LatestTime, 1)))
            val logSize = consumer.getOffsetsBefore(request).partitionErrorAndOffsets(topicAndPartition).offsets.head
            println("%20s%d".format("Log size = ", logSize))
            println("%20s%,d (%,.2fG)".format("= ", logSize, logSize / math.pow(1024, 3)))

            val lag = logSize - offset
            println("%20s%d".format("Consumer lag = ", lag))
            println("%20s%,d (%,.2fG)".format("= ", lag, lag / math.pow(1024, 3)))
            println()
          case None => // ignore
        }
      case _ =>
        error("Could not parse broker/partition pair %s".format(bidPid))
    }
  }

  private def processTopic(zkClient: ZkClient, group: String, topic: String) {
    val bidsPids = ZkUtils.getChildrenParentMayNotExist(
      zkClient, "/consumers/%s/offsets/%s".format(group, topic)).toList
    bidsPids.sorted.foreach {
      bidPid => processPartition(zkClient, group, topic, bidPid)
    }
  }

  private def printBrokerInfo() {
    println("BROKER INFO")
    for ((bid, consumerOpt) <- consumerMap)
      consumerOpt match {
        case Some(consumer) =>
          println("%s -> %s:%d".format(bid, consumer.host, consumer.port))
        case None => // ignore
      }
  }

  def main(args: Array[String]) {
    val parser = new OptionParser()

    val zkConnectOpt = parser.accepts("zkconnect", "ZooKeeper connect string.").
            withRequiredArg().defaultsTo("localhost:2181").ofType(classOf[String]);
    val topicsOpt = parser.accepts("topic",
            "Comma-separated list of consumer topics (all topics if absent).").
            withRequiredArg().ofType(classOf[String])
    val groupOpt = parser.accepts("group", "Consumer group.").
            withRequiredArg().ofType(classOf[String])
    parser.accepts("help", "Print this message.")

    val options = parser.parse(args : _*)

    if (options.has("help")) {
       parser.printHelpOn(System.out)
       System.exit(0)
    }

    for (opt <- List(groupOpt))
      if (!options.has(opt)) {
        System.err.println("Missing required argument: %s".format(opt))
        parser.printHelpOn(System.err)
        System.exit(1)
      }

    val zkConnect = options.valueOf(zkConnectOpt)
    val group = options.valueOf(groupOpt)
    val topics = if (options.has(topicsOpt)) Some(options.valueOf(topicsOpt))
      else None


    var zkClient: ZkClient = null
    try {
      zkClient = new ZkClient(zkConnect, 30000, 30000, ZKStringSerializer)

      val topicList = topics match {
        case Some(x) => x.split(",").view.toList
        case None => ZkUtils.getChildren(
          zkClient, "/consumers/%s/offsets".format(group)).toList
      }

      debug("zkConnect = %s; topics = %s; group = %s".format(
        zkConnect, topicList.toString(), group))

      topicList.sorted.foreach {
        topic => processTopic(zkClient, group, topic)
      }

      printBrokerInfo()
    }
    finally {
      for (consumerOpt <- consumerMap.values) {
        consumerOpt match {
          case Some(consumer) => consumer.close()
          case None => // ignore
        }
      }
      if (zkClient != null)
        zkClient.close()
    }
  }
}
<commit_msg>ConsumerOffsetChecker does not work with 0.8; kafka-685; patched by <NAME>; reviewed by <NAME><commit_after>/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package kafka.tools


import joptsimple._
import org.I0Itec.zkclient.ZkClient
import kafka.utils.{ZkUtils, ZKStringSerializer, Logging}
import kafka.consumer.SimpleConsumer
import kafka.api.{PartitionOffsetRequestInfo, OffsetRequest}
import kafka.common.TopicAndPartition
import scala.collection._


object ConsumerOffsetChecker extends Logging {

  private val consumerMap: mutable.Map[Int, Option[SimpleConsumer]] = mutable.Map()

  private val BrokerIpPattern = """^([^:]+):(\d+).*$""".r
  // e.g., 127.0.0.1:9092:9999 (with JMX port)

  private def getConsumer(zkClient: ZkClient, bid: Int): Option[SimpleConsumer] = {
    val brokerInfo = ZkUtils.readDataMaybeNull(zkClient, "/brokers/ids/%s".format(bid))._1
    val consumer = brokerInfo match {
      case Some(BrokerIpPattern(ip, port)) =>
        Some(new SimpleConsumer(ip, port.toInt, 10000, 100000, "ConsumerOffsetChecker"))
      case _ =>
        error("Could not parse broker info %s with regex %s".format(brokerInfo, BrokerIpPattern.toString()))
        None
    }
    consumer
  }

  private def processPartition(zkClient: ZkClient,
                               group: String, topic: String, pid: Int) {
    val offset = ZkUtils.readData(zkClient, "/consumers/%s/offsets/%s/%s".
            format(group, topic, pid))._1.toLong
    val owner = ZkUtils.readDataMaybeNull(zkClient, "/consumers/%s/owners/%s/%s".
            format(group, topic, pid))._1

    ZkUtils.getLeaderForPartition(zkClient, topic, pid) match {
      case Some(bid) =>
        val consumerOpt = consumerMap.getOrElseUpdate(bid, getConsumer(zkClient, bid))
        consumerOpt match {
          case Some(consumer) =>
            val topicAndPartition = TopicAndPartition(topic, pid)
            val request =
              OffsetRequest(immutable.Map(topicAndPartition -> PartitionOffsetRequestInfo(OffsetRequest.LatestTime, 1)))
            val logSize = consumer.getOffsetsBefore(request).partitionErrorAndOffsets(topicAndPartition).offsets.head

            val lag = logSize - offset
            println("%-15s %-30s %-3s %-15s %-15s %-15s %s".format(group, topic, pid, offset, logSize, lag,
              owner match {case Some(ownerStr) => ownerStr case None => "none"}))
            consumer.close()
          case None => // ignore
        }
      case None =>
        error("No broker for partition %s - %s".format(topic, pid))
    }
  }

  private def processTopic(zkClient: ZkClient, group: String, topic: String) {
    val pidMap = ZkUtils.getPartitionsForTopics(zkClient, Seq(topic))
    pidMap.get(topic) match {
      case Some(pids) =>
        pids.sorted.foreach {
          pid => processPartition(zkClient, group, topic, pid)
        }
      case None => // ignore
    }
  }

  private def printBrokerInfo() {
    println("BROKER INFO")
    for ((bid, consumerOpt) <- consumerMap)
      consumerOpt match {
        case Some(consumer) =>
          println("%s -> %s:%d".format(bid, consumer.host, consumer.port))
        case None => // ignore
      }
  }

  def main(args: Array[String]) {
    val parser = new OptionParser()

    val zkConnectOpt = parser.accepts("zkconnect", "ZooKeeper connect string.").
            withRequiredArg().defaultsTo("localhost:2181").ofType(classOf[String]);
    val topicsOpt = parser.accepts("topic",
            "Comma-separated list of consumer topics (all topics if absent).").
            withRequiredArg().ofType(classOf[String])
    val groupOpt = parser.accepts("group", "Consumer group.").
            withRequiredArg().ofType(classOf[String])
    parser.accepts("broker-info", "Print broker info")
    parser.accepts("help", "Print this message.")

    val options = parser.parse(args : _*)

    if (options.has("help")) {
       parser.printHelpOn(System.out)
       System.exit(0)
    }

    for (opt <- List(groupOpt))
      if (!options.has(opt)) {
        System.err.println("Missing required argument: %s".format(opt))
        parser.printHelpOn(System.err)
        System.exit(1)
      }

    val zkConnect = options.valueOf(zkConnectOpt)
    val group = options.valueOf(groupOpt)
    val topics = if (options.has(topicsOpt)) Some(options.valueOf(topicsOpt))
      else None


    var zkClient: ZkClient = null
    try {
      zkClient = new ZkClient(zkConnect, 30000, 30000, ZKStringSerializer)

      val topicList = topics match {
        case Some(x) => x.split(",").view.toList
        case None => ZkUtils.getChildren(
          zkClient, "/consumers/%s/offsets".format(group)).toList
      }

      debug("zkConnect = %s; topics = %s; group = %s".format(
        zkConnect, topicList.toString(), group))

      println("%-15s %-30s %-3s %-15s %-15s %-15s %s".format("Group", "Topic", "Pid", "Offset", "logSize", "Lag", "Owner"))
      topicList.sorted.foreach {
        topic => processTopic(zkClient, group, topic)
      }

      if (options.has("broker-info"))
        printBrokerInfo();

    }
    finally {
      for (consumerOpt <- consumerMap.values) {
        consumerOpt match {
          case Some(consumer) => consumer.close()
          case None => // ignore
        }
      }
      if (zkClient != null)
        zkClient.close()
    }
  }
}

{'fraction_non_alphanumeric': 0.1272801302931596, 'fraction_numerical': 0.015553745928338763, 'mean_word_length': 2.9514157014157014, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 41, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 3, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7389521', 'n_tokens_mistral': 3814, 'n_tokens_neox': 3603, 'n_words': 922}