[LI-HOTFIX] Add dynamic maintenance broker config

TICKET = KAFKA-8527
LI_DESCRIPTION =
When a broker is masked as maintenance broker, Kafka will not
assign partitions of new topics to the broker if partition assignment is done by
kafka brokers.
1) add dynamic broker config to support maintenance brokers
2) add logic to update maintenance broker info in controller
3) support maintenance broker in Kafka API
4) support maintenance broker in AdminZkClient
5) add logic to rearrange replica assignment for new topics for maintenance brokers
5) test cases

RB=1484649
G=Kafka-Code-Reviews
R=agencer, amendhek, jkoshy, zahuang
A=zahuang

EXIT_CRITERIA = TICKET [KAFKA-8527]

Author: xiowu0

core/src/main/scala/kafka/controller/KafkaController.scala

@@ -95,6 +95,7 @@ class KafkaController(val config: KafkaConfig,
                       threadNamePrefix: Option[String] = None)
   extends ControllerEventProcessor with Logging with KafkaMetricsGroup {
 
+  val adminZkClient = new AdminZkClient(zkClient)
   this.logIdent = s"[Controller id=${config.brokerId}] "
 
   @volatile private var brokerInfo = initialBrokerInfo
@@ -180,6 +181,13 @@ class KafkaController(val config: KafkaConfig,
     }
   )
 
+  newGauge(
+    "MaintenanceBrokerCount",
+    new Gauge[Int] {
+      def value: Int = if (isActive) config.getMaintenanceBrokerList.size else 0
+    }
+  )
+
   /**
    * Returns true if this broker is the current controller.
    */
@@ -695,6 +703,7 @@ class KafkaController(val config: KafkaConfig,
     controllerContext.setLiveBrokerAndEpochs(curBrokerAndEpochs)
     info(s"Initialized broker epochs cache: ${controllerContext.liveBrokerIdAndEpochs}")
     controllerContext.allTopics = zkClient.getAllTopicsInCluster.toSet
+    rearrangePartitionReplicaAssignmentForNewTopics(controllerContext.allTopics)
     registerPartitionModificationsHandlers(controllerContext.allTopics.toSeq)
     getReplicaAssignmentPolicyCompliant(controllerContext.allTopics.toSet).foreach {
       case (topicPartition, assignedReplicas) => controllerContext.updatePartitionReplicaAssignment(topicPartition, assignedReplicas)
@@ -770,6 +779,37 @@ class KafkaController(val config: KafkaConfig,
     }
   }
 
+  // Rearrange partition and replica assignment for new topics that get assigned to
+  // maintenance brokers that do not take new partitions
+  private def rearrangePartitionReplicaAssignmentForNewTopics(topics: Set[String]) {
+    try {
+      val noNewPartitionBrokerIds = config.getMaintenanceBrokerList
+      if (noNewPartitionBrokerIds.nonEmpty) {
+        val newTopics = zkClient.getPartitionNodeNonExistsTopics(topics.toSet)
+        val newTopicsToBeArranged = zkClient.getPartitionAssignmentForTopics(newTopics).filter {
+          case (_, partitionMap) =>
+            partitionMap.exists {
+              case (_, assignedReplicas) =>
+                assignedReplicas.intersect(noNewPartitionBrokerIds).nonEmpty
+            }
+        }
+        newTopicsToBeArranged.foreach {
+          case (topic, partitionMap) =>
+            val numPartitions = partitionMap.size
+            val numReplica = partitionMap.head._2.size
+            val brokers = controllerContext.liveOrShuttingDownBrokers.map { b => kafka.admin.BrokerMetadata(b.id, b.rack) }.toSeq
+
+            val replicaAssignment = adminZkClient.assignReplicasToAvailableBrokers(brokers, noNewPartitionBrokerIds.toSet, numPartitions, numReplica)
+            adminZkClient.writeTopicPartitionAssignment(topic, replicaAssignment, true)
+            info(s"Rearrange partition and replica assignment for topic [$topic]")
+        }
+      }
+    } catch {
+      case e =>
+        error("Error during rearranging partition and replica assignment for new topics for maintenance brokers :" + e.getMessage)
+    }
+  }
+
   private def moveReassignedPartitionLeaderIfRequired(topicPartition: TopicPartition,
                                                       reassignedPartitionContext: ReassignedPartitionsContext) {
     val reassignedReplicas = reassignedPartitionContext.newReplicas
@@ -1320,6 +1360,7 @@ class KafkaController(val config: KafkaConfig,
     val newTopics = topics -- controllerContext.allTopics
     val deletedTopics = controllerContext.allTopics -- topics
     controllerContext.allTopics = topics
+    rearrangePartitionReplicaAssignmentForNewTopics(newTopics)
 
     registerPartitionModificationsHandlers(newTopics.toSeq)
     val addedPartitionReplicaAssignment = getReplicaAssignmentPolicyCompliant(newTopics)
@@ -1593,11 +1634,12 @@ class KafkaController(val config: KafkaConfig,
 
     val replicationFactor = config.defaultReplicationFactor
     val brokers = controllerContext.liveOrShuttingDownBrokers.map { sb => kafka.admin.BrokerMetadata(sb.id, sb.rack) }.toSeq
+    val noNewPartitionBrokerIds = config.getMaintenanceBrokerList.toSet
 
     topicsReplicaAssignment.foreach{
       case(topic, partitionAssignment) => {
         val numPartitions = partitionAssignment.size
-        val assignment = AdminUtils.assignReplicasToBrokers(brokers, numPartitions, replicationFactor)
+        val assignment = adminZkClient.assignReplicasToAvailableBrokers(brokers, noNewPartitionBrokerIds, numPartitions, replicationFactor)
           .map{ case(partition, replicas) => {
             (new TopicPartition(topic, partition), replicas)
           }}.toMap

core/src/main/scala/kafka/server/AdminManager.scala

@@ -18,7 +18,7 @@ package kafka.server
 
 import java.util.{Collections, Properties}
 
-import kafka.admin.{AdminOperationException, AdminUtils}
+import kafka.admin.AdminOperationException
 import kafka.common.TopicAlreadyMarkedForDeletionException
 import kafka.log.LogConfig
 import kafka.metrics.KafkaMetricsGroup
@@ -96,7 +96,7 @@ class AdminManager(val config: KafkaConfig,
             "Both cannot be used at the same time.")
         }
         val assignments = if (topic.assignments().isEmpty) {
-          AdminUtils.assignReplicasToBrokers(brokers, topic.numPartitions, topic.replicationFactor)
+          adminZkClient.assignReplicasToAvailableBrokers(brokers, config.getMaintenanceBrokerList.toSet, topic.numPartitions, topic.replicationFactor)
         } else {
           val assignments = new mutable.HashMap[Int, Seq[Int]]
           // Note: we don't check that replicaAssignment contains unknown brokers - unlike in add-partitions case,
@@ -276,7 +276,8 @@ class AdminManager(val config: KafkaConfig,
         }
 
         val updatedReplicaAssignment = adminZkClient.addPartitions(topic, existingAssignment, allBrokers,
-          newPartition.totalCount, reassignment, validateOnly = validateOnly)
+          newPartition.totalCount, reassignment, validateOnly = validateOnly,
+          noNewPartitionBrokerIds = config.getMaintenanceBrokerList.toSet)
         CreatePartitionsMetadata(topic, updatedReplicaAssignment, ApiError.NONE)
       } catch {
         case e: AdminOperationException =>
@@ -543,7 +544,10 @@ class AdminManager(val config: KafkaConfig,
   }
 
   private def configType(name: String, synonyms: List[String]): ConfigDef.Type = {
-    val configType = config.typeOf(name)
+    var configType = config.typeOf(name)
+    if (configType == null)
+      configType = DynamicConfig.Broker.typeOf(name)
+
     if (configType != null)
       configType
     else

core/src/main/scala/kafka/server/DynamicBrokerConfig.scala

@@ -130,7 +130,10 @@ object DynamicBrokerConfig {
     checkInvalidProps(securityConfigsWithoutListenerPrefix(props),
       "These security configs can be dynamically updated only per-listener using the listener prefix")
     validateConfigTypes(props)
-    if (!perBrokerConfig) {
+    if (perBrokerConfig) {
+      checkInvalidProps(clusterLevelConfigs(props),
+        "Cannot update these configs at per broker level, broker id must not be specified")
+    } else {
       checkInvalidProps(perBrokerConfigs(props),
         "Cannot update these configs at default cluster level, broker id must be specified")
     }
@@ -147,6 +150,11 @@ object DynamicBrokerConfig {
     configNames.intersect(PerBrokerConfigs) ++ configNames.filter(perBrokerListenerConfig)
   }
 
+  private def clusterLevelConfigs(props: Properties): Set[String] = {
+    val configNames = props.asScala.keySet
+    configNames.intersect(DynamicConfig.Broker.ClusterLevelConfigs)
+  }
+
   private def nonDynamicConfigs(props: Properties): Set[String] = {
     props.asScala.keySet.intersect(DynamicConfig.Broker.nonDynamicProps)
   }
@@ -311,6 +319,11 @@ class DynamicBrokerConfig(private val kafkaConfig: KafkaConfig) extends Logging
       }
   }
 
+  private[server] def getMaintenanceBrokerList: Seq[Int] = CoreUtils.inReadLock(lock) {
+    DynamicConfig.Broker.getMaintenanceBrokerListFromString(dynamicDefaultConfigs.getOrElse(DynamicConfig.Broker.MaintenanceBrokerListProp,
+      DynamicConfig.Broker.DefaultMaintenanceBrokerList).toString)
+  }
+
   private def maybeCreatePasswordEncoder(secret: Option[Password]): Option[PasswordEncoder] = {
    secret.map { secret =>
       new PasswordEncoder(secret,

core/src/main/scala/kafka/server/DynamicConfig.scala

@@ -21,10 +21,11 @@ import java.util.Properties
 
 import kafka.log.LogConfig
 import kafka.security.CredentialProvider
-import org.apache.kafka.common.config.ConfigDef
+import org.apache.kafka.common.config.{ConfigDef, ConfigException}
 import org.apache.kafka.common.config.ConfigDef.Importance._
 import org.apache.kafka.common.config.ConfigDef.Range._
 import org.apache.kafka.common.config.ConfigDef.Type._
+import org.apache.kafka.common.config.ConfigDef.Validator
 
 import scala.collection.JavaConverters._
 
@@ -39,9 +40,11 @@ object DynamicConfig {
     val LeaderReplicationThrottledRateProp = "leader.replication.throttled.rate"
     val FollowerReplicationThrottledRateProp = "follower.replication.throttled.rate"
     val ReplicaAlterLogDirsIoMaxBytesPerSecondProp = "replica.alter.log.dirs.io.max.bytes.per.second"
+    val MaintenanceBrokerListProp = "maintenance.broker.list"
 
     //Defaults
     val DefaultReplicationThrottledRate = ReplicationQuotaManagerConfig.QuotaBytesPerSecondDefault
+    val DefaultMaintenanceBrokerList: String = ""
 
     //Documentation
     val LeaderReplicationThrottledRateDoc = "A long representing the upper bound (bytes/sec) on replication traffic for leaders enumerated in the " +
@@ -52,19 +55,45 @@ object DynamicConfig {
       s"limit be kept above 1MB/s for accurate behaviour."
     val ReplicaAlterLogDirsIoMaxBytesPerSecondDoc = "A long representing the upper bound (bytes/sec) on disk IO used for moving replica between log directories on the same broker. " +
       s"This property can be only set dynamically. It is suggested that the limit be kept above 1MB/s for accurate behaviour."
+    val MaintenanceBrokerListDoc = "A list containing maintenance broker Ids, separated by comma"
 
     //Definitions
     private val brokerConfigDef = new ConfigDef()
       //round minimum value down, to make it easier for users.
       .define(LeaderReplicationThrottledRateProp, LONG, DefaultReplicationThrottledRate, atLeast(0), MEDIUM, LeaderReplicationThrottledRateDoc)
       .define(FollowerReplicationThrottledRateProp, LONG, DefaultReplicationThrottledRate, atLeast(0), MEDIUM, FollowerReplicationThrottledRateDoc)
       .define(ReplicaAlterLogDirsIoMaxBytesPerSecondProp, LONG, DefaultReplicationThrottledRate, atLeast(0), MEDIUM, ReplicaAlterLogDirsIoMaxBytesPerSecondDoc)
+      .define(MaintenanceBrokerListProp, STRING, DefaultMaintenanceBrokerList, MaintenanceBrokerListValidator, MEDIUM, MaintenanceBrokerListDoc)
     DynamicBrokerConfig.addDynamicConfigs(brokerConfigDef)
     val nonDynamicProps = KafkaConfig.configNames.toSet -- brokerConfigDef.names.asScala
 
+    //cluster level only configs
+    val ClusterLevelConfigs = Set(MaintenanceBrokerListProp)
+    def typeOf(key: String): ConfigDef.Type = {
+      val configKey: ConfigDef.ConfigKey = brokerConfigDef.configKeys().get(key)
+      if (configKey == null)
+        null
+      else
+        configKey.`type`
+    }
+
     def names = brokerConfigDef.names
 
     def validate(props: Properties) = DynamicConfig.validate(brokerConfigDef, props, customPropsAllowed = true)
+
+    def getMaintenanceBrokerListFromString(brokerListStr: String): Seq[Int] = {
+      brokerListStr.split(",").map(_.trim).filter(_.nonEmpty).map(_.toInt)
+    }
+
+    object MaintenanceBrokerListValidator extends Validator {
+      override def ensureValid(name: String, value: Any): Unit = {
+        try {
+          getMaintenanceBrokerListFromString(value.toString)
+        } catch {
+          case e: NumberFormatException => throw new ConfigException(name, value.toString, e.getMessage)
+        }
+      }
+    }
   }
 
   object Client {

core/src/main/scala/kafka/server/KafkaConfig.scala

@@ -1403,6 +1403,10 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean, dynamicConfigO
     millis
   }
 
+  def getMaintenanceBrokerList: Seq[Int] = {
+    dynamicConfig.getMaintenanceBrokerList
+  }
+
   private def getMap(propName: String, propValue: String): Map[String, String] = {
     try {
       CoreUtils.parseCsvMap(propValue)

core/src/main/scala/kafka/zk/AdminZkClient.scala

@@ -52,7 +52,8 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
                   topicConfig: Properties = new Properties,
                   rackAwareMode: RackAwareMode = RackAwareMode.Enforced) {
     val brokerMetadatas = getBrokerMetadatas(rackAwareMode)
-    val replicaAssignment = AdminUtils.assignReplicasToBrokers(brokerMetadatas, partitions, replicationFactor)
+    val noNewPartitionBrokerIds = getMaintenanceBrokerList()
+    val replicaAssignment = assignReplicasToAvailableBrokers(brokerMetadatas, noNewPartitionBrokerIds.toSet, partitions, replicationFactor)
     createTopicWithAssignment(topic, topicConfig, replicaAssignment)
   }
 
@@ -80,6 +81,16 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
     brokerMetadatas.sortBy(_.id)
   }
 
+  /**
+    * fetch maintenance broker list from zk
+    */
+  def getMaintenanceBrokerList(): Seq[Int] = {
+    val maintenanceBrokerConfig = fetchEntityConfig(ConfigType.Broker, ConfigEntityName.Default)
+      .getProperty(DynamicConfig.Broker.MaintenanceBrokerListProp, DynamicConfig.Broker.DefaultMaintenanceBrokerList)
+
+    DynamicConfig.Broker.getMaintenanceBrokerListFromString(maintenanceBrokerConfig)
+  }
+
   def createTopicWithAssignment(topic: String,
                                 config: Properties,
                                 partitionReplicaAssignment: Map[Int, Seq[Int]]): Unit = {
@@ -134,7 +145,7 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
     LogConfig.validate(config)
   }
 
-  private def writeTopicPartitionAssignment(topic: String, replicaAssignment: Map[Int, Seq[Int]], isUpdate: Boolean) {
+  def writeTopicPartitionAssignment(topic: String, replicaAssignment: Map[Int, Seq[Int]], isUpdate: Boolean) {
     try {
       val assignment = replicaAssignment.map { case (partitionId, replicas) => (new TopicPartition(topic,partitionId), replicas) }.toMap
 
@@ -168,6 +179,39 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
     }
   }
 
+  /**
+    * Assign replicas to brokers that take new partitions.
+    * If the number of replicationFactor is greater than the number of brokers that take new partitions,
+    * all brokers are used for assignment.
+    * @param brokerMetadatas
+    * @param noNewPartitionBrokerIds
+    * @param nPartitions
+    * @param replicationFactor
+    * @param fixedStartIndex
+    * @param startPartitionId
+    * @return
+    */
+  def assignReplicasToAvailableBrokers(brokerMetadatas: Seq[BrokerMetadata],
+    noNewPartitionBrokerIds: Set[Int],
+    nPartitions: Int,
+    replicationFactor: Int,
+    fixedStartIndex: Int = -1,
+    startPartitionId: Int = -1): Map[Int, Seq[Int]] = {
+
+    val availableBrokerMetadata = brokerMetadatas.filter {
+      brokerMetadata =>
+        if (noNewPartitionBrokerIds.contains(brokerMetadata.id)) false
+        else true
+    }
+
+    if (replicationFactor > availableBrokerMetadata.size) {
+      info(s"Using all brokers for replica assignment since replicationFactor[$replicationFactor] " +
+        s"is larger than the number of nonMaintenanceBroker[${availableBrokerMetadata.size}]")
+      AdminUtils.assignReplicasToBrokers(brokerMetadatas, nPartitions, replicationFactor, fixedStartIndex, startPartitionId)
+    } else
+      AdminUtils.assignReplicasToBrokers(availableBrokerMetadata, nPartitions, replicationFactor, fixedStartIndex, startPartitionId)
+  }
+
   /**
   * Add partitions to existing topic with optional replica assignment
   *
@@ -177,14 +221,16 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
   * @param numPartitions Number of partitions to be set
   * @param replicaAssignment Manual replica assignment, or none
   * @param validateOnly If true, validate the parameters without actually adding the partitions
+  * @param noNewPartitionBrokerIds Brokers that do not take new partitions
   * @return the updated replica assignment
   */
   def addPartitions(topic: String,
                     existingAssignment: Map[Int, Seq[Int]],
                     allBrokers: Seq[BrokerMetadata],
                     numPartitions: Int = 1,
                     replicaAssignment: Option[Map[Int, Seq[Int]]] = None,
-                    validateOnly: Boolean = false): Map[Int, Seq[Int]] = {
+                    validateOnly: Boolean = false,
+                    noNewPartitionBrokerIds: Set[Int] = Set.empty[Int]): Map[Int, Seq[Int]] = {
     val existingAssignmentPartition0 = existingAssignment.getOrElse(0,
       throw new AdminOperationException(
         s"Unexpected existing replica assignment for topic '$topic', partition id 0 is missing. " +
@@ -204,8 +250,8 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
 
     val proposedAssignmentForNewPartitions = replicaAssignment.getOrElse {
       val startIndex = math.max(0, allBrokers.indexWhere(_.id >= existingAssignmentPartition0.head))
-      AdminUtils.assignReplicasToBrokers(allBrokers, partitionsToAdd, existingAssignmentPartition0.size,
-        startIndex, existingAssignment.size)
+      assignReplicasToAvailableBrokers(allBrokers, noNewPartitionBrokerIds, partitionsToAdd,
+        existingAssignmentPartition0.size, startIndex, existingAssignment.size)
     }
 
     val proposedAssignment = existingAssignment ++ proposedAssignmentForNewPartitions

core/src/main/scala/kafka/zk/KafkaZkClient.scala

@@ -602,6 +602,23 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
     }.toMap
   }
 
+  def getPartitionNodeNonExistsTopics(topics: Set[String]): Set[String] = {
+    val existsRequests = topics.map(topic => ExistsRequest(TopicPartitionsZNode.path(topic), ctx = Some(topic)))
+    val existsResponses = retryRequestsUntilConnected(existsRequests.toSeq)
+    val newTopics = scala.collection.mutable.Set.empty[String]
+
+    existsResponses.foreach {
+      existsResponse =>
+        val topic = existsResponse.ctx.get.asInstanceOf[String]
+        existsResponse.resultCode match {
+          case Code.OK =>
+          case Code.NONODE => newTopics.add(topic)
+          case _ => throw existsResponse.resultException.get
+        }
+    }
+    newTopics.toSet
+  }
+
   /**
    * Gets the partition numbers for the given topics
    * @param topics the topics whose partitions we wish to get.