Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 71 additions & 17 deletions core/src/main/scala/kafka/controller/ControllerContext.scala
Original file line number Diff line number Diff line change
Expand Up @@ -31,14 +31,46 @@ class ControllerContext {
var epoch: Int = KafkaController.InitialControllerEpoch - 1
var epochZkVersion: Int = KafkaController.InitialControllerEpochZkVersion - 1
var allTopics: Set[String] = Set.empty
var partitionReplicaAssignment: mutable.Map[TopicPartition, Seq[Int]] = mutable.Map.empty
var partitionLeadershipInfo: mutable.Map[TopicPartition, LeaderIsrAndControllerEpoch] = mutable.Map.empty
private var partitionReplicaAssignmentUnderlying: mutable.Map[String, mutable.Map[Int, Seq[Int]]] = mutable.Map.empty
val partitionLeadershipInfo: mutable.Map[TopicPartition, LeaderIsrAndControllerEpoch] = mutable.Map.empty
val partitionsBeingReassigned: mutable.Map[TopicPartition, ReassignedPartitionsContext] = mutable.Map.empty
val replicasOnOfflineDirs: mutable.Map[Int, Set[TopicPartition]] = mutable.Map.empty

private var liveBrokersUnderlying: Set[Broker] = Set.empty
private var liveBrokerIdsUnderlying: Set[Int] = Set.empty

def partitionReplicaAssignment(topicPartition: TopicPartition) : Seq[Int] = {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we remove the space before : Seq[Int] and some other places?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed.

partitionReplicaAssignmentUnderlying.getOrElse(topicPartition.topic, mutable.Map.empty)
.getOrElse(topicPartition.partition, Seq.empty)
}

def clearTopicsState(): Unit = {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could this be private?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed.

allTopics = Set.empty
partitionReplicaAssignmentUnderlying.clear()
partitionLeadershipInfo.clear()
partitionsBeingReassigned.clear()
replicasOnOfflineDirs.clear()
}

def updatePartitionReplicaAssignment(topicPartition: TopicPartition, newReplicas: Seq[Int]) : Unit = {
partitionReplicaAssignmentUnderlying.getOrElseUpdate(topicPartition.topic, mutable.Map.empty)
.put(topicPartition.partition, newReplicas)
}

def partitionReplicaAssignmentForTopic(topic : String) : Map[TopicPartition, Seq[Int]] = {
partitionReplicaAssignmentUnderlying.getOrElse(topic, Map.empty).map {
case (partition, replicas) => (new TopicPartition(topic, partition), replicas)
}.toMap
}

def allPartitions : Set[TopicPartition] = {
partitionReplicaAssignmentUnderlying.flatMap {
case (topic, topicReplicaAssignment) => topicReplicaAssignment.map {
case (partition, _) => new TopicPartition(topic, partition)
}
}.toSet
}

// setter
def liveBrokers_=(brokers: Set[Broker]) {
liveBrokersUnderlying = brokers
Expand All @@ -53,8 +85,12 @@ class ControllerContext {
def liveOrShuttingDownBrokers = liveBrokersUnderlying

def partitionsOnBroker(brokerId: Int): Set[TopicPartition] = {
partitionReplicaAssignment.collect {
case (topicPartition, replicas) if replicas.contains(brokerId) => topicPartition
partitionReplicaAssignmentUnderlying.flatMap {
case (topic, topicReplicaAssignment) => topicReplicaAssignment.filter {
case (_, replicas) => replicas.contains(brokerId)
}.map {
case (partition, _) => new TopicPartition(topic, partition)
}
}.toSet
}

Expand All @@ -68,22 +104,26 @@ class ControllerContext {

def replicasOnBrokers(brokerIds: Set[Int]): Set[PartitionAndReplica] = {
brokerIds.flatMap { brokerId =>
partitionReplicaAssignment.collect { case (topicPartition, replicas) if replicas.contains(brokerId) =>
PartitionAndReplica(topicPartition, brokerId)
partitionReplicaAssignmentUnderlying.flatMap {
case (topic, topicReplicaAssignment) => topicReplicaAssignment.collect {
case (partition, replicas) if replicas.contains(brokerId) =>
PartitionAndReplica(new TopicPartition(topic, partition), brokerId)
}
}
}.toSet
}
}

def replicasForTopic(topic: String): Set[PartitionAndReplica] = {
partitionReplicaAssignment
.filter { case (topicPartition, _) => topicPartition.topic == topic }
.flatMap { case (topicPartition, replicas) =>
replicas.map(PartitionAndReplica(topicPartition, _))
}.toSet
partitionReplicaAssignmentUnderlying.getOrElse(topic, mutable.Map.empty).flatMap {
case (partition, replicas) => replicas.map(r => PartitionAndReplica(new TopicPartition(topic, partition), r))
}.toSet
}

def partitionsForTopic(topic: String): collection.Set[TopicPartition] =
partitionReplicaAssignment.keySet.filter(topicPartition => topicPartition.topic == topic)
def partitionsForTopic(topic: String): collection.Set[TopicPartition] = {
partitionReplicaAssignmentUnderlying.getOrElse(topic, mutable.Map.empty).map {
case (partition, _) => new TopicPartition(topic, partition)
}.toSet
}

def allLiveReplicas(): Set[PartitionAndReplica] = {
replicasOnBrokers(liveBrokerIds).filter { partitionAndReplica =>
Expand All @@ -98,10 +138,24 @@ class ControllerContext {
}
}

def resetContext() : Unit = {
if (controllerChannelManager != null) {
controllerChannelManager.shutdown()
controllerChannelManager = null
}
shuttingDownBrokerIds.clear()
epoch = 0
epochZkVersion = 0
clearTopicsState()
liveBrokers = Set.empty
}

def removeTopic(topic: String) = {
partitionLeadershipInfo = partitionLeadershipInfo.filter { case (topicPartition, _) => topicPartition.topic != topic }
partitionReplicaAssignment = partitionReplicaAssignment.filter { case (topicPartition, _) => topicPartition.topic != topic }
allTopics -= topic
partitionReplicaAssignmentUnderlying.remove(topic)
partitionLeadershipInfo.foreach {
case (topicPartition, _) if topicPartition.topic == topic => partitionLeadershipInfo.remove(topicPartition)
case _ =>
}
}

}
117 changes: 56 additions & 61 deletions core/src/main/scala/kafka/controller/KafkaController.scala
Original file line number Diff line number Diff line change
Expand Up @@ -309,7 +309,7 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
replicaStateMachine.shutdown()
zkClient.unregisterZNodeChildChangeHandler(brokerChangeHandler.path)

resetControllerContext()
controllerContext.resetContext()

info("Resigned")
}
Expand Down Expand Up @@ -569,28 +569,28 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
}
val newReplicas = reassignedPartitionContext.newReplicas
val topic = tp.topic
controllerContext.partitionReplicaAssignment.get(tp) match {
case Some(assignedReplicas) =>
if (assignedReplicas == newReplicas) {
info(s"Partition $tp to be reassigned is already assigned to replicas " +
s"${newReplicas.mkString(",")}. Ignoring request for partition reassignment.")
removePartitionFromReassignedPartitions(tp)
} else {
try {
info(s"Handling reassignment of partition $tp to new replicas ${newReplicas.mkString(",")}")
// first register ISR change listener
reassignedPartitionContext.registerReassignIsrChangeHandler(zkClient)
// mark topic ineligible for deletion for the partitions being reassigned
topicDeletionManager.markTopicIneligibleForDeletion(Set(topic))
onPartitionReassignment(tp, reassignedPartitionContext)
} catch {
case e: Throwable =>
error(s"Error completing reassignment of partition $tp", e)
// remove the partition from the admin path to unblock the admin client
removePartitionFromReassignedPartitions(tp)
}
val assignedReplicas = controllerContext.partitionReplicaAssignment(tp)
if (assignedReplicas.nonEmpty) {
if (assignedReplicas == newReplicas) {
info(s"Partition $tp to be reassigned is already assigned to replicas " +
s"${newReplicas.mkString(",")}. Ignoring request for partition reassignment.")
removePartitionFromReassignedPartitions(tp)
} else {
try {
info(s"Handling reassignment of partition $tp to new replicas ${newReplicas.mkString(",")}")
// first register ISR change listener
reassignedPartitionContext.registerReassignIsrChangeHandler(zkClient)
// mark topic ineligible for deletion for the partitions being reassigned
topicDeletionManager.markTopicIneligibleForDeletion(Set(topic))
onPartitionReassignment(tp, reassignedPartitionContext)
} catch {
case e: Throwable =>
error(s"Error completing reassignment of partition $tp", e)
// remove the partition from the admin path to unblock the admin client
removePartitionFromReassignedPartitions(tp)
}
case None =>
}
} else {
error(s"Ignoring request to reassign partition $tp that doesn't exist.")
removePartitionFromReassignedPartitions(tp)
}
Expand Down Expand Up @@ -643,8 +643,10 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
controllerContext.liveBrokers = zkClient.getAllBrokersInCluster.toSet
controllerContext.allTopics = zkClient.getAllTopicsInCluster.toSet
registerPartitionModificationsHandlers(controllerContext.allTopics.toSeq)
controllerContext.partitionReplicaAssignment = mutable.Map.empty ++ zkClient.getReplicaAssignmentForTopics(controllerContext.allTopics.toSet)
controllerContext.partitionLeadershipInfo = new mutable.HashMap[TopicPartition, LeaderIsrAndControllerEpoch]
zkClient.getReplicaAssignmentForTopics(controllerContext.allTopics.toSet).foreach {
case (topicPartition, assignedReplicas) => controllerContext.updatePartitionReplicaAssignment(topicPartition, assignedReplicas)
}
controllerContext.partitionLeadershipInfo.clear()
controllerContext.shuttingDownBrokerIds = mutable.Set.empty[Int]
// register broker modifications handlers
registerBrokerModificationsHandler(controllerContext.liveBrokers.map(_.id))
Expand All @@ -662,10 +664,10 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
val partitionsUndergoingPreferredReplicaElection = zkClient.getPreferredReplicaElection
// check if they are already completed or topic was deleted
val partitionsThatCompletedPreferredReplicaElection = partitionsUndergoingPreferredReplicaElection.filter { partition =>
val replicasOpt = controllerContext.partitionReplicaAssignment.get(partition)
val topicDeleted = replicasOpt.isEmpty
val replicas = controllerContext.partitionReplicaAssignment(partition)
val topicDeleted = replicas.isEmpty
val successful =
if (!topicDeleted) controllerContext.partitionLeadershipInfo(partition).leaderAndIsr.leader == replicasOpt.get.head else false
if (!topicDeleted) controllerContext.partitionLeadershipInfo(partition).leaderAndIsr.leader == replicas.head else false
successful || topicDeleted
}
val pendingPreferredReplicaElectionsIgnoringTopicDeletion = partitionsUndergoingPreferredReplicaElection -- partitionsThatCompletedPreferredReplicaElection
Expand All @@ -678,21 +680,6 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
pendingPreferredReplicaElections
}

private def resetControllerContext(): Unit = {
if (controllerContext.controllerChannelManager != null) {
controllerContext.controllerChannelManager.shutdown()
controllerContext.controllerChannelManager = null
}
controllerContext.shuttingDownBrokerIds.clear()
controllerContext.epoch = 0
controllerContext.epochZkVersion = 0
controllerContext.allTopics = Set.empty
controllerContext.partitionReplicaAssignment.clear()
controllerContext.partitionLeadershipInfo.clear()
controllerContext.partitionsBeingReassigned.clear()
controllerContext.liveBrokers = Set.empty
}

private def initializePartitionReassignment() {
// read the partitions being reassigned from zookeeper path /admin/reassign_partitions
val partitionsBeingReassigned = zkClient.getPartitionReassignment
Expand All @@ -706,9 +693,10 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti

private def fetchTopicDeletionsInProgress(): (Set[String], Set[String]) = {
val topicsToBeDeleted = zkClient.getTopicDeletions.toSet
val topicsWithOfflineReplicas = controllerContext.partitionReplicaAssignment.filter { case (partition, replicas) =>
replicas.exists(r => !controllerContext.isReplicaOnline(r, partition))
}.keySet.map(_.topic)
val topicsWithOfflineReplicas = controllerContext.allTopics.filter { topic => {
val replicasForTopic = controllerContext.replicasForTopic(topic)
replicasForTopic.exists(r => !controllerContext.isReplicaOnline(r.replica, r.topicPartition))
}}
val topicsForWhichPartitionReassignmentIsInProgress = controllerContext.partitionsBeingReassigned.keySet.map(_.topic)
val topicsIneligibleForDeletion = topicsWithOfflineReplicas | topicsForWhichPartitionReassignmentIsInProgress
info(s"List of topics to be deleted: ${topicsToBeDeleted.mkString(",")}")
Expand All @@ -722,7 +710,7 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
controllerContext.controllerChannelManager.startup()
}

private def updateLeaderAndIsrCache(partitions: Seq[TopicPartition] = controllerContext.partitionReplicaAssignment.keys.toSeq) {
private def updateLeaderAndIsrCache(partitions: Seq[TopicPartition] = controllerContext.allPartitions.toSeq) {
val leaderIsrAndControllerEpochs = zkClient.getTopicPartitionStates(partitions)
leaderIsrAndControllerEpochs.foreach { case (partition, leaderIsrAndControllerEpoch) =>
controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
Expand All @@ -742,7 +730,7 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
// change the assigned replica list to just the reassigned replicas in the cache so it gets sent out on the LeaderAndIsr
// request to the current or new leader. This will prevent it from adding the old replicas to the ISR
val oldAndNewReplicas = controllerContext.partitionReplicaAssignment(topicPartition)
controllerContext.partitionReplicaAssignment.put(topicPartition, reassignedReplicas)
controllerContext.updatePartitionReplicaAssignment(topicPartition, reassignedReplicas)
if (!reassignedPartitionContext.newReplicas.contains(currentLeader)) {
info(s"Leader $currentLeader for partition $topicPartition being reassigned, " +
s"is not in the new list of replicas ${reassignedReplicas.mkString(",")}. Re-electing leader")
Expand Down Expand Up @@ -778,14 +766,13 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti

private def updateAssignedReplicasForPartition(partition: TopicPartition,
replicas: Seq[Int]) {
val partitionsAndReplicasForThisTopic = controllerContext.partitionReplicaAssignment.filter(_._1.topic == partition.topic)
partitionsAndReplicasForThisTopic.put(partition, replicas)
val setDataResponse = zkClient.setTopicAssignmentRaw(partition.topic, partitionsAndReplicasForThisTopic.toMap)
controllerContext.updatePartitionReplicaAssignment(partition, replicas)
val setDataResponse = zkClient.setTopicAssignmentRaw(partition.topic, controllerContext.partitionReplicaAssignmentForTopic(partition.topic))
setDataResponse.resultCode match {
case Code.OK =>
info(s"Updated assigned replicas for partition $partition being reassigned to ${replicas.mkString(",")}")
// update the assigned replica list after a successful zookeeper write
controllerContext.partitionReplicaAssignment.put(partition, replicas)
controllerContext.updatePartitionReplicaAssignment(partition, replicas)
case Code.NONODE => throw new IllegalStateException(s"Topic ${partition.topic} doesn't exist")
case _ => throw new KafkaException(setDataResponse.resultException.get)
}
Expand Down Expand Up @@ -971,9 +958,12 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
private def checkAndTriggerAutoLeaderRebalance(): Unit = {
trace("Checking need to trigger auto leader balancing")
val preferredReplicasForTopicsByBrokers: Map[Int, Map[TopicPartition, Seq[Int]]] =
controllerContext.partitionReplicaAssignment.filterNot { case (tp, _) =>
topicDeletionManager.isTopicQueuedUpForDeletion(tp.topic)
}.groupBy { case (_, assignedReplicas) => assignedReplicas.head }
controllerContext.allPartitions.filterNot {
tp => topicDeletionManager.isTopicQueuedUpForDeletion(tp.topic)
}.map { tp =>
(tp, controllerContext.partitionReplicaAssignment(tp) )
}.toMap.groupBy { case (_, assignedReplicas) => assignedReplicas.head }

debug(s"Preferred replicas by broker $preferredReplicasForTopicsByBrokers")

// for each broker, check if a preferred replica election needs to be triggered
Expand Down Expand Up @@ -1155,7 +1145,8 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
if (!isActive) {
0
} else {
controllerContext.partitionReplicaAssignment.count { case (topicPartition, replicas) =>
controllerContext.allPartitions.count { topicPartition =>
val replicas = controllerContext.partitionReplicaAssignment(topicPartition)
val preferredReplica = replicas.head
val leadershipInfo = controllerContext.partitionLeadershipInfo.get(topicPartition)
leadershipInfo.map(_.leaderAndIsr.leader != preferredReplica).getOrElse(false) &&
Expand Down Expand Up @@ -1279,9 +1270,10 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti

registerPartitionModificationsHandlers(newTopics.toSeq)
val addedPartitionReplicaAssignment = zkClient.getReplicaAssignmentForTopics(newTopics)
controllerContext.partitionReplicaAssignment = controllerContext.partitionReplicaAssignment.filter(p =>
!deletedTopics.contains(p._1.topic))
controllerContext.partitionReplicaAssignment ++= addedPartitionReplicaAssignment
deletedTopics.foreach(controllerContext.removeTopic)
addedPartitionReplicaAssignment.foreach {
case (topicAndPartition, newReplicas) => controllerContext.updatePartitionReplicaAssignment(topicAndPartition, newReplicas)
}
info(s"New topics: [$newTopics], deleted topics: [$deletedTopics], new partition replica assignment " +
s"[$addedPartitionReplicaAssignment]")
if (addedPartitionReplicaAssignment.nonEmpty)
Expand Down Expand Up @@ -1311,15 +1303,18 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
override def process(): Unit = {
if (!isActive) return
val partitionReplicaAssignment = zkClient.getReplicaAssignmentForTopics(immutable.Set(topic))
val partitionsToBeAdded = partitionReplicaAssignment.filter(p =>
!controllerContext.partitionReplicaAssignment.contains(p._1))
val partitionsToBeAdded = partitionReplicaAssignment.filter { case (topicPartition, _) =>
controllerContext.partitionReplicaAssignment(topicPartition).isEmpty
}
if (topicDeletionManager.isTopicQueuedUpForDeletion(topic))
error(s"Skipping adding partitions ${partitionsToBeAdded.map(_._1.partition).mkString(",")} for topic $topic " +
"since it is currently being deleted")
else {
if (partitionsToBeAdded.nonEmpty) {
info(s"New partitions to be added $partitionsToBeAdded")
controllerContext.partitionReplicaAssignment ++= partitionsToBeAdded
partitionsToBeAdded.foreach { case (topicPartition, assignedReplicas) =>
controllerContext.updatePartitionReplicaAssignment(topicPartition, assignedReplicas)
}
onNewPartitionCreation(partitionsToBeAdded.keySet)
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ class PartitionStateMachine(config: KafkaConfig,
* zookeeper
*/
private def initializePartitionState() {
for (topicPartition <- controllerContext.partitionReplicaAssignment.keys) {
for (topicPartition <- controllerContext.allPartitions) {
// check if leader and isr path exists for partition. If not, then it is in NEW state
controllerContext.partitionLeadershipInfo.get(topicPartition) match {
case Some(currentLeaderIsrAndEpoch) =>
Expand Down
Loading