KAFKA-2639: Refactoring of ZkUtils

core/src/main/scala/kafka/admin/ConfigCommand.scala

@@ -21,7 +21,7 @@ import joptsimple._
 import java.util.Properties
 import kafka.log.LogConfig
 import kafka.server.ConfigType
-import kafka.utils.{ZkUtils, CommandLineUtils}
+import kafka.utils.{ToolsUtils, ZkUtils, CommandLineUtils}
 import org.I0Itec.zkclient.ZkClient
 import scala.collection._
 import scala.collection.JavaConversions._
@@ -42,52 +42,55 @@ object ConfigCommand {
 
     opts.checkArgs()
 
-    val zkClient = ZkUtils.createZkClient(opts.options.valueOf(opts.zkConnectOpt), 30000, 30000)
+    val zkUtils = ZkUtils.create(opts.options.valueOf(opts.zkConnectOpt),
+                                 30000,
+                                 30000,
+                                 ToolsUtils.isSecure(System.getProperty("java.security.auth.login.config")))
 
     try {
       if (opts.options.has(opts.alterOpt))
-        alterConfig(zkClient, opts)
+        alterConfig(zkUtils, opts)
       else if (opts.options.has(opts.describeOpt))
-        describeConfig(zkClient, opts)
+        describeConfig(zkUtils, opts)
     } catch {
       case e: Throwable =>
         println("Error while executing topic command " + e.getMessage)
         println(Utils.stackTrace(e))
     } finally {
-      zkClient.close()
+      zkUtils.close()
     }
   }
 
-  private def alterConfig(zkClient: ZkClient, opts: ConfigCommandOptions) {
+  private def alterConfig(zkUtils: ZkUtils, opts: ConfigCommandOptions) {
     val configsToBeAdded = parseConfigsToBeAdded(opts)
     val configsToBeDeleted = parseConfigsToBeDeleted(opts)
     val entityType = opts.options.valueOf(opts.entityType)
     val entityName = opts.options.valueOf(opts.entityName)
 
     // compile the final set of configs
-    val configs = AdminUtils.fetchEntityConfig(zkClient, entityType, entityName)
+    val configs = AdminUtils.fetchEntityConfig(zkUtils, entityType, entityName)
     configs.putAll(configsToBeAdded)
     configsToBeDeleted.foreach(config => configs.remove(config))
 
     if (entityType.equals(ConfigType.Topic)) {
-      AdminUtils.changeTopicConfig(zkClient, entityName, configs)
+      AdminUtils.changeTopicConfig(zkUtils, entityName, configs)
       println("Updated config for topic: \"%s\".".format(entityName))
     } else {
-      AdminUtils.changeClientIdConfig(zkClient, entityName, configs)
+      AdminUtils.changeClientIdConfig(zkUtils, entityName, configs)
       println("Updated config for clientId: \"%s\".".format(entityName))
     }
   }
 
-  private def describeConfig(zkClient: ZkClient, opts: ConfigCommandOptions) {
+  private def describeConfig(zkUtils: ZkUtils, opts: ConfigCommandOptions) {
     val entityType = opts.options.valueOf(opts.entityType)
     val entityNames: Seq[String] =
       if (opts.options.has(opts.entityName))
         Seq(opts.options.valueOf(opts.entityName))
       else
-        ZkUtils.getAllEntitiesWithConfig(zkClient, entityType)
+        zkUtils.getAllEntitiesWithConfig(entityType)
 
     for (entityName <- entityNames) {
-      val configs = AdminUtils.fetchEntityConfig(zkClient, entityType, entityName)
+      val configs = AdminUtils.fetchEntityConfig(zkUtils, entityType, entityName)
       println("Configs for %s:%s are %s"
                       .format(entityType, entityName, configs.map(kv => kv._1 + "=" + kv._2).mkString(",")))
     }

core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala

@@ -48,57 +48,60 @@ object ConsumerGroupCommand {
 
     opts.checkArgs()
 
-    val zkClient = ZkUtils.createZkClient(opts.options.valueOf(opts.zkConnectOpt), 30000, 30000)
+    val zkUtils = ZkUtils.create(opts.options.valueOf(opts.zkConnectOpt), 
+                                 30000,
+                                 30000,
+                                 ToolsUtils.isSecure(System.getProperty("java.security.auth.login.config")))
 
     try {
       if (opts.options.has(opts.listOpt))
-        list(zkClient)
+        list(zkUtils)
       else if (opts.options.has(opts.describeOpt))
-        describe(zkClient, opts)
+        describe(zkUtils, opts)
       else if (opts.options.has(opts.deleteOpt))
-        delete(zkClient, opts)
+        delete(zkUtils, opts)
     } catch {
       case e: Throwable =>
         println("Error while executing consumer group command " + e.getMessage)
         println(Utils.stackTrace(e))
     } finally {
-      zkClient.close()
+      zkUtils.close()
     }
   }
 
-  def list(zkClient: ZkClient) {
-    ZkUtils.getConsumerGroups(zkClient).foreach(println)
+  def list(zkUtils: ZkUtils) {
+    zkUtils.getConsumerGroups().foreach(println)
   }
 
-  def describe(zkClient: ZkClient, opts: ConsumerGroupCommandOptions) {
+  def describe(zkUtils: ZkUtils, opts: ConsumerGroupCommandOptions) {
     val configs = parseConfigs(opts)
     val channelSocketTimeoutMs = configs.getProperty("channelSocketTimeoutMs", "600").toInt
     val channelRetryBackoffMs = configs.getProperty("channelRetryBackoffMsOpt", "300").toInt
     val group = opts.options.valueOf(opts.groupOpt)
-    val topics = ZkUtils.getTopicsByConsumerGroup(zkClient, group)
+    val topics = zkUtils.getTopicsByConsumerGroup(group)
     if (topics.isEmpty) {
       println("No topic available for consumer group provided")
     }
-    topics.foreach(topic => describeTopic(zkClient, group, topic, channelSocketTimeoutMs, channelRetryBackoffMs))
+    topics.foreach(topic => describeTopic(zkUtils, group, topic, channelSocketTimeoutMs, channelRetryBackoffMs))
   }
 
-  def delete(zkClient: ZkClient, opts: ConsumerGroupCommandOptions) {
+  def delete(zkUtils: ZkUtils, opts: ConsumerGroupCommandOptions) {
     if (opts.options.has(opts.groupOpt) && opts.options.has(opts.topicOpt)) {
-      deleteForTopic(zkClient, opts)
+      deleteForTopic(zkUtils, opts)
     }
     else if (opts.options.has(opts.groupOpt)) {
-      deleteForGroup(zkClient, opts)
+      deleteForGroup(zkUtils, opts)
     }
     else if (opts.options.has(opts.topicOpt)) {
-      deleteAllForTopic(zkClient, opts)
+      deleteAllForTopic(zkUtils, opts)
     }
   }
 
-  private def deleteForGroup(zkClient: ZkClient, opts: ConsumerGroupCommandOptions) {
+  private def deleteForGroup(zkUtils: ZkUtils, opts: ConsumerGroupCommandOptions) {
     val groups = opts.options.valuesOf(opts.groupOpt)
     groups.foreach { group =>
       try {
-        if (AdminUtils.deleteConsumerGroupInZK(zkClient, group))
+        if (AdminUtils.deleteConsumerGroupInZK(zkUtils, group))
           println("Deleted all consumer group information for group %s in zookeeper.".format(group))
         else
           println("Delete for group %s failed because its consumers are still active.".format(group))
@@ -110,13 +113,13 @@ object ConsumerGroupCommand {
     }
   }
 
-  private def deleteForTopic(zkClient: ZkClient, opts: ConsumerGroupCommandOptions) {
+  private def deleteForTopic(zkUtils: ZkUtils, opts: ConsumerGroupCommandOptions) {
     val groups = opts.options.valuesOf(opts.groupOpt)
     val topic = opts.options.valueOf(opts.topicOpt)
     Topic.validate(topic)
     groups.foreach { group =>
       try {
-        if (AdminUtils.deleteConsumerGroupInfoForTopicInZK(zkClient, group, topic))
+        if (AdminUtils.deleteConsumerGroupInfoForTopicInZK(zkUtils, group, topic))
           println("Deleted consumer group information for group %s topic %s in zookeeper.".format(group, topic))
         else
           println("Delete for group %s topic %s failed because its consumers are still active.".format(group, topic))
@@ -128,10 +131,10 @@ object ConsumerGroupCommand {
     }
   }
 
-  private def deleteAllForTopic(zkClient: ZkClient, opts: ConsumerGroupCommandOptions) {
+  private def deleteAllForTopic(zkUtils: ZkUtils, opts: ConsumerGroupCommandOptions) {
     val topic = opts.options.valueOf(opts.topicOpt)
     Topic.validate(topic)
-    AdminUtils.deleteAllConsumerGroupInfoForTopicInZK(zkClient, topic)
+    AdminUtils.deleteAllConsumerGroupInfoForTopicInZK(zkUtils, topic)
     println("Deleted consumer group information for all inactive consumer groups for topic %s in zookeeper.".format(topic))
   }
 
@@ -144,35 +147,35 @@ object ConsumerGroupCommand {
     props
   }
 
-  private def describeTopic(zkClient: ZkClient,
+  private def describeTopic(zkUtils: ZkUtils,
                             group: String,
                             topic: String,
                             channelSocketTimeoutMs: Int,
                             channelRetryBackoffMs: Int) {
-    val topicPartitions = getTopicPartitions(zkClient, topic)
-    val partitionOffsets = getPartitionOffsets(zkClient, group, topicPartitions, channelSocketTimeoutMs, channelRetryBackoffMs)
+    val topicPartitions = getTopicPartitions(zkUtils, topic)
+    val partitionOffsets = getPartitionOffsets(zkUtils, group, topicPartitions, channelSocketTimeoutMs, channelRetryBackoffMs)
     println("%s, %s, %s, %s, %s, %s, %s"
       .format("GROUP", "TOPIC", "PARTITION", "CURRENT OFFSET", "LOG END OFFSET", "LAG", "OWNER"))
     topicPartitions
       .sortBy { case topicPartition => topicPartition.partition }
       .foreach { topicPartition =>
-      describePartition(zkClient, group, topicPartition.topic, topicPartition.partition, partitionOffsets.get(topicPartition))
+      describePartition(zkUtils, group, topicPartition.topic, topicPartition.partition, partitionOffsets.get(topicPartition))
     }
   }
 
-  private def getTopicPartitions(zkClient: ZkClient, topic: String) = {
-    val topicPartitionMap = ZkUtils.getPartitionsForTopics(zkClient, Seq(topic))
+  private def getTopicPartitions(zkUtils: ZkUtils, topic: String) = {
+    val topicPartitionMap = zkUtils.getPartitionsForTopics(Seq(topic))
     val partitions = topicPartitionMap.getOrElse(topic, Seq.empty)
     partitions.map(TopicAndPartition(topic, _))
   }
 
-  private def getPartitionOffsets(zkClient: ZkClient,
+  private def getPartitionOffsets(zkUtils: ZkUtils,
                                   group: String,
                                   topicPartitions: Seq[TopicAndPartition],
                                   channelSocketTimeoutMs: Int,
                                   channelRetryBackoffMs: Int): Map[TopicAndPartition, Long] = {
     val offsetMap = mutable.Map[TopicAndPartition, Long]()
-    val channel = ClientUtils.channelToOffsetManager(group, zkClient, channelSocketTimeoutMs, channelRetryBackoffMs)
+    val channel = ClientUtils.channelToOffsetManager(group, zkUtils, channelSocketTimeoutMs, channelRetryBackoffMs)
     channel.send(OffsetFetchRequest(group, topicPartitions))
     val offsetFetchResponse = OffsetFetchResponse.readFrom(channel.receive().payload())
 
@@ -182,7 +185,7 @@ object ConsumerGroupCommand {
         // this group may not have migrated off zookeeper for offsets storage (we don't expose the dual-commit option in this tool
         // (meaning the lag may be off until all the consumers in the group have the same setting for offsets storage)
         try {
-          val offset = ZkUtils.readData(zkClient, topicDirs.consumerOffsetDir + "/" + topicAndPartition.partition)._1.toLong
+          val offset = zkUtils.readData(topicDirs.consumerOffsetDir + "/" + topicAndPartition.partition)._1.toLong
           offsetMap.put(topicAndPartition, offset)
         } catch {
           case z: ZkNoNodeException =>
@@ -200,20 +203,20 @@ object ConsumerGroupCommand {
     offsetMap.toMap
   }
 
-  private def describePartition(zkClient: ZkClient,
+  private def describePartition(zkUtils: ZkUtils,
                                 group: String,
                                 topic: String,
                                 partition: Int,
                                 offsetOpt: Option[Long]) {
     val topicAndPartition = TopicAndPartition(topic, partition)
     val groupDirs = new ZKGroupTopicDirs(group, topic)
-    val owner = ZkUtils.readDataMaybeNull(zkClient, groupDirs.consumerOwnerDir + "/" + partition)._1
-    ZkUtils.getLeaderForPartition(zkClient, topic, partition) match {
+    val owner = zkUtils.readDataMaybeNull(groupDirs.consumerOwnerDir + "/" + partition)._1
+    zkUtils.getLeaderForPartition(topic, partition) match {
       case Some(-1) =>
         println("%s, %s, %s, %s, %s, %s, %s"
           .format(group, topic, partition, offsetOpt.getOrElse("unknown"), "unknown", "unknown", owner.getOrElse("none")))
       case Some(brokerId) =>
-        val consumerOpt = getConsumer(zkClient, brokerId)
+        val consumerOpt = getConsumer(zkUtils, brokerId)
         consumerOpt match {
           case Some(consumer) =>
             val request =
@@ -231,9 +234,9 @@ object ConsumerGroupCommand {
     }
   }
 
-  private def getConsumer(zkClient: ZkClient, brokerId: Int): Option[SimpleConsumer] = {
+  private def getConsumer(zkUtils: ZkUtils, brokerId: Int): Option[SimpleConsumer] = {
     try {
-      ZkUtils.readDataMaybeNull(zkClient, ZkUtils.BrokerIdsPath + "/" + brokerId)._1 match {
+      zkUtils.readDataMaybeNull(ZkUtils.BrokerIdsPath + "/" + brokerId)._1 match {
         case Some(brokerInfoString) =>
           Json.parseFull(brokerInfoString) match {
             case Some(m) =>

core/src/main/scala/kafka/admin/PreferredReplicaLeaderElectionCommand.scala

@@ -51,15 +51,19 @@ object PreferredReplicaLeaderElectionCommand extends Logging {
 
     val zkConnect = options.valueOf(zkConnectOpt)
     var zkClient: ZkClient = null
-
+    var zkUtils: ZkUtils = null
     try {
       zkClient = ZkUtils.createZkClient(zkConnect, 30000, 30000)
+      zkUtils = ZkUtils.create(zkConnect, 
+                               30000,
+                               30000,
+                               ToolsUtils.isSecure(System.getProperty("java.security.auth.login.config")))
       val partitionsForPreferredReplicaElection =
         if (!options.has(jsonFileOpt))
-          ZkUtils.getAllPartitions(zkClient)
+          zkUtils.getAllPartitions()
         else
           parsePreferredReplicaElectionData(Utils.readFileAsString(options.valueOf(jsonFileOpt)))
-      val preferredReplicaElectionCommand = new PreferredReplicaLeaderElectionCommand(zkClient, partitionsForPreferredReplicaElection)
+      val preferredReplicaElectionCommand = new PreferredReplicaLeaderElectionCommand(zkUtils, partitionsForPreferredReplicaElection)
 
       preferredReplicaElectionCommand.moveLeaderToPreferredReplica()
       println("Successfully started preferred replica election for partitions %s".format(partitionsForPreferredReplicaElection))
@@ -95,39 +99,39 @@ object PreferredReplicaLeaderElectionCommand extends Logging {
     }
   }
 
-  def writePreferredReplicaElectionData(zkClient: ZkClient,
+  def writePreferredReplicaElectionData(zkUtils: ZkUtils,
                                         partitionsUndergoingPreferredReplicaElection: scala.collection.Set[TopicAndPartition]) {
     val zkPath = ZkUtils.PreferredReplicaLeaderElectionPath
     val partitionsList = partitionsUndergoingPreferredReplicaElection.map(e => Map("topic" -> e.topic, "partition" -> e.partition))
     val jsonData = Json.encode(Map("version" -> 1, "partitions" -> partitionsList))
     try {
-      ZkUtils.createPersistentPath(zkClient, zkPath, jsonData)
+      zkUtils.createPersistentPath(zkPath, jsonData)
       info("Created preferred replica election path with %s".format(jsonData))
     } catch {
       case nee: ZkNodeExistsException =>
         val partitionsUndergoingPreferredReplicaElection =
-          PreferredReplicaLeaderElectionCommand.parsePreferredReplicaElectionData(ZkUtils.readData(zkClient, zkPath)._1)
+          PreferredReplicaLeaderElectionCommand.parsePreferredReplicaElectionData(zkUtils.readData(zkPath)._1)
         throw new AdminOperationException("Preferred replica leader election currently in progress for " +
           "%s. Aborting operation".format(partitionsUndergoingPreferredReplicaElection))
       case e2: Throwable => throw new AdminOperationException(e2.toString)
     }
   }
 }
 
-class PreferredReplicaLeaderElectionCommand(zkClient: ZkClient, partitions: scala.collection.Set[TopicAndPartition])
+class PreferredReplicaLeaderElectionCommand(zkUtils: ZkUtils, partitions: scala.collection.Set[TopicAndPartition])
   extends Logging {
   def moveLeaderToPreferredReplica() = {
     try {
-      val validPartitions = partitions.filter(p => validatePartition(zkClient, p.topic, p.partition))
-      PreferredReplicaLeaderElectionCommand.writePreferredReplicaElectionData(zkClient, validPartitions)
+      val validPartitions = partitions.filter(p => validatePartition(zkUtils, p.topic, p.partition))
+      PreferredReplicaLeaderElectionCommand.writePreferredReplicaElectionData(zkUtils, validPartitions)
     } catch {
       case e: Throwable => throw new AdminCommandFailedException("Admin command failed", e)
     }
   }
 
-  def validatePartition(zkClient: ZkClient, topic: String, partition: Int): Boolean = {
+  def validatePartition(zkUtils: ZkUtils, topic: String, partition: Int): Boolean = {
     // check if partition exists
-    val partitionsOpt = ZkUtils.getPartitionsForTopics(zkClient, List(topic)).get(topic)
+    val partitionsOpt = zkUtils.getPartitionsForTopics(List(topic)).get(topic)
     partitionsOpt match {
       case Some(partitions) =>
         if(partitions.contains(partition)) {