Skip to content

MINOR: Improved error handling in ZK migration #13372

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cmccabe merged 15 commits into from
Mar 16, 2023
Merged

MINOR: Improved error handling in ZK migration #13372

Show file tree
Hide file tree
Changes from 8 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
160b20a
Catch KeeperException in ZkMigrationClient, add unit test
mumrah Mar 9, 2023
d139892
Improve WAIT_FOR_BROKERS performance
mumrah Mar 9, 2023
e8dded8
checkstyle
mumrah Mar 9, 2023
c2316cd
Create second exception type
mumrah Mar 9, 2023
d5c5631
Use mutable set
mumrah Mar 9, 2023
9ffb908
fix flaky test failure
mumrah Mar 10, 2023
b2ed202
Merge remote-tracking branch 'origin/trunk' into minor-better-zk-migr…
mumrah Mar 13, 2023
fde1737
Update to new method name
mumrah Mar 13, 2023
86bb0e4
PR feedback
mumrah Mar 14, 2023
53deae1
Merge remote-tracking branch 'origin/trunk' into minor-better-zk-migr…
mumrah Mar 14, 2023
fa4f10f
checkstyle
mumrah Mar 14, 2023
95d5675
Fix Scala 2.12 build
mumrah Mar 15, 2023
230a917
Add timeout to ZkIntegrationTest
mumrah Mar 15, 2023
20345d8
Merge remote-tracking branch 'origin/trunk' into minor-better-zk-migr…
mumrah Mar 16, 2023
1d4f60e
Fix KRaftMigrationDriverTest
mumrah Mar 16, 2023
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
135 changes: 84 additions & 51 deletions core/src/main/scala/kafka/zk/ZkMigrationClient.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,11 +28,10 @@ import org.apache.kafka.common.metadata.ClientQuotaRecord.EntityData
import org.apache.kafka.common.metadata._
import org.apache.kafka.common.quota.ClientQuotaEntity
import org.apache.kafka.common.{TopicPartition, Uuid}
import org.apache.kafka.image.{MetadataDelta, MetadataImage}
import org.apache.kafka.metadata.{LeaderRecoveryState, PartitionRegistration}
import org.apache.kafka.metadata.migration.{MigrationClient, ZkMigrationLeadershipState}
import org.apache.kafka.metadata.migration.{MigrationClient, MigrationClientAuthException, MigrationClientException, ZkMigrationLeadershipState}
import org.apache.kafka.server.common.{ApiMessageAndVersion, ProducerIdsBlock}
import org.apache.zookeeper.KeeperException.Code
import org.apache.zookeeper.KeeperException.{AuthFailedException, Code, NoAuthException, SessionClosedRequireAuthException}
import org.apache.zookeeper.{CreateMode, KeeperException}

import java.util
Expand All @@ -47,24 +46,44 @@ import scala.jdk.CollectionConverters._
*/
class ZkMigrationClient(zkClient: KafkaZkClient) extends MigrationClient with Logging {
Comment thread
mumrah marked this conversation as resolved.
Outdated

override def getOrCreateMigrationRecoveryState(initialState: ZkMigrationLeadershipState): ZkMigrationLeadershipState = {
zkClient.createTopLevelPaths()
zkClient.getOrCreateMigrationState(initialState)
@throws(classOf[MigrationClientException])
def wrapZkException[T](fn: => T): T = {
Comment thread
mumrah marked this conversation as resolved.
Outdated
try {
fn
} catch {
case e @ (_: AuthFailedException | _: NoAuthException | _: SessionClosedRequireAuthException) =>
// We don't expect authentication errors to be recoverable, so treat them differently
throw new MigrationClientAuthException(e)
case e: KeeperException => throw new MigrationClientException(e)
}
}

override def setMigrationRecoveryState(state: ZkMigrationLeadershipState): ZkMigrationLeadershipState = {
override def getOrCreateMigrationRecoveryState(
initialState: ZkMigrationLeadershipState
): ZkMigrationLeadershipState = wrapZkException {
zkClient.createTopLevelPaths()
zkClient.getOrCreateMigrationState(initialState)
}

override def setMigrationRecoveryState(
state: ZkMigrationLeadershipState
): ZkMigrationLeadershipState = wrapZkException {
zkClient.updateMigrationState(state)
}

override def claimControllerLeadership(state: ZkMigrationLeadershipState): ZkMigrationLeadershipState = {
override def claimControllerLeadership(
state: ZkMigrationLeadershipState
): ZkMigrationLeadershipState = wrapZkException {
zkClient.tryRegisterKRaftControllerAsActiveController(state.kraftControllerId(), state.kraftControllerEpoch()) match {
case SuccessfulRegistrationResult(controllerEpoch, controllerEpochZkVersion) =>
state.withZkController(controllerEpoch, controllerEpochZkVersion)
case FailedRegistrationResult() => state.withUnknownZkController()
}
}

override def releaseControllerLeadership(state: ZkMigrationLeadershipState): ZkMigrationLeadershipState = {
override def releaseControllerLeadership(
state: ZkMigrationLeadershipState
): ZkMigrationLeadershipState = wrapZkException {
try {
zkClient.deleteController(state.zkControllerEpochZkVersion())
state.withUnknownZkController()
Expand All @@ -73,7 +92,7 @@ class ZkMigrationClient(zkClient: KafkaZkClient) extends MigrationClient with Lo
// If the controller moved, no need to release
state.withUnknownZkController()
case t: Throwable =>
throw new RuntimeException("Could not release controller leadership due to underlying error", t)
throw new MigrationClientException("Could not release controller leadership due to underlying error", t)
}
}

Expand Down Expand Up @@ -211,19 +230,21 @@ class ZkMigrationClient(zkClient: KafkaZkClient) extends MigrationClient with Lo
}
}

override def readAllMetadata(batchConsumer: Consumer[util.List[ApiMessageAndVersion]],
brokerIdConsumer: Consumer[Integer]): Unit = {
override def readAllMetadata(
batchConsumer: Consumer[util.List[ApiMessageAndVersion]],
brokerIdConsumer: Consumer[Integer]
): Unit = wrapZkException {
Comment thread
mumrah marked this conversation as resolved.
Outdated
migrateTopics(batchConsumer, brokerIdConsumer)
migrateBrokerConfigs(batchConsumer)
migrateClientQuotas(batchConsumer)
migrateProducerId(batchConsumer)
}

override def readBrokerIds(): util.Set[Integer] = {
override def readBrokerIds(): util.Set[Integer] = wrapZkException {
zkClient.getSortedBrokerList.map(Integer.valueOf).toSet.asJava
}

override def readBrokerIdsFromTopicAssignments(): util.Set[Integer] = {
override def readBrokerIdsFromTopicAssignments(): util.Set[Integer] = wrapZkException {
val topics = zkClient.getAllTopicsInCluster()
val replicaAssignmentAndTopicIds = zkClient.getReplicaAssignmentAndTopicIdForTopics(topics)
val brokersWithAssignments = new util.HashSet[Integer]()
Expand All @@ -235,10 +256,12 @@ class ZkMigrationClient(zkClient: KafkaZkClient) extends MigrationClient with Lo
brokersWithAssignments
}

override def createTopic(topicName: String,
topicId: Uuid,
partitions: util.Map[Integer, PartitionRegistration],
state: ZkMigrationLeadershipState): ZkMigrationLeadershipState = {
override def createTopic(
topicName: String,
topicId: Uuid,
partitions: util.Map[Integer, PartitionRegistration],
state: ZkMigrationLeadershipState
): ZkMigrationLeadershipState = wrapZkException {
val assignments = partitions.asScala.map { case (partitionId, partition) =>
new TopicPartition(topicName, partitionId) ->
ReplicaAssignment(partition.replicas, partition.addingReplicas, partition.removingReplicas)
Expand Down Expand Up @@ -280,18 +303,22 @@ class ZkMigrationClient(zkClient: KafkaZkClient) extends MigrationClient with Lo
state.withMigrationZkVersion(migrationZkVersion)
} else {
// not ok
throw new RuntimeException(s"Failed to create or update topic $topicName. ZK operation had results $resultCodes")
throw new MigrationClientException(s"Failed to create or update topic $topicName. ZK operation had results $resultCodes")
}
}

private def createTopicPartition(topicPartition: TopicPartition): CreateRequest = {
private def createTopicPartition(
topicPartition: TopicPartition
): CreateRequest = wrapZkException {
val path = TopicPartitionZNode.path(topicPartition)
CreateRequest(path, null, zkClient.defaultAcls(path), CreateMode.PERSISTENT, Some(topicPartition))
}

private def partitionStatePathAndData(topicPartition: TopicPartition,
partitionRegistration: PartitionRegistration,
controllerEpoch: Int): (String, Array[Byte]) = {
private def partitionStatePathAndData(
topicPartition: TopicPartition,
partitionRegistration: PartitionRegistration,
controllerEpoch: Int
): (String, Array[Byte]) = {
val path = TopicPartitionStateZNode.path(topicPartition)
val data = TopicPartitionStateZNode.encode(LeaderIsrAndControllerEpoch(new LeaderAndIsr(
partitionRegistration.leader,
Expand All @@ -302,22 +329,28 @@ class ZkMigrationClient(zkClient: KafkaZkClient) extends MigrationClient with Lo
(path, data)
}

private def createTopicPartitionState(topicPartition: TopicPartition,
partitionRegistration: PartitionRegistration,
controllerEpoch: Int): CreateRequest = {
private def createTopicPartitionState(
topicPartition: TopicPartition,
partitionRegistration: PartitionRegistration,
controllerEpoch: Int
): CreateRequest = {
val (path, data) = partitionStatePathAndData(topicPartition, partitionRegistration, controllerEpoch)
CreateRequest(path, data, zkClient.defaultAcls(path), CreateMode.PERSISTENT, Some(topicPartition))
}

private def updateTopicPartitionState(topicPartition: TopicPartition,
partitionRegistration: PartitionRegistration,
controllerEpoch: Int): SetDataRequest = {
private def updateTopicPartitionState(
topicPartition: TopicPartition,
partitionRegistration: PartitionRegistration,
controllerEpoch: Int
): SetDataRequest = {
val (path, data) = partitionStatePathAndData(topicPartition, partitionRegistration, controllerEpoch)
SetDataRequest(path, data, ZkVersion.MatchAnyVersion, Some(topicPartition))
}

override def updateTopicPartitions(topicPartitions: util.Map[String, util.Map[Integer, PartitionRegistration]],
state: ZkMigrationLeadershipState): ZkMigrationLeadershipState = {
override def updateTopicPartitions(
topicPartitions: util.Map[String, util.Map[Integer, PartitionRegistration]],
state: ZkMigrationLeadershipState
): ZkMigrationLeadershipState = wrapZkException {
val requests = topicPartitions.asScala.flatMap { case (topicName, partitionRegistrations) =>
partitionRegistrations.asScala.flatMap { case (partitionId, partitionRegistration) =>
val topicPartition = new TopicPartition(topicName, partitionId)
Expand All @@ -332,18 +365,20 @@ class ZkMigrationClient(zkClient: KafkaZkClient) extends MigrationClient with Lo
if (resultCodes.forall { case (_, code) => code.equals(Code.OK) } ) {
state.withMigrationZkVersion(migrationZkVersion)
} else {
throw new RuntimeException(s"Failed to update partition states: $topicPartitions. ZK transaction had results $resultCodes")
throw new MigrationClientException(s"Failed to update partition states: $topicPartitions. ZK transaction had results $resultCodes")
}
}
}

// Try to update an entity config and the migration state. If NoNode is encountered, it probably means we
// need to recursively create the parent ZNode. In this case, return None.
def tryWriteEntityConfig(entityType: String,
path: String,
props: Properties,
create: Boolean,
state: ZkMigrationLeadershipState): Option[ZkMigrationLeadershipState] = {
def tryWriteEntityConfig(
entityType: String,
path: String,
props: Properties,
create: Boolean,
state: ZkMigrationLeadershipState
): Option[ZkMigrationLeadershipState] = wrapZkException {
val configData = ConfigEntityZNode.encode(props)

val requests = if (create) {
Expand All @@ -366,8 +401,7 @@ class ZkMigrationClient(zkClient: KafkaZkClient) extends MigrationClient with Lo
entity: util.Map[String, String],
quotas: util.Map[String, java.lang.Double],
state: ZkMigrationLeadershipState
): ZkMigrationLeadershipState = {

): ZkMigrationLeadershipState = wrapZkException {
val entityMap = entity.asScala
val hasUser = entityMap.contains(ClientQuotaEntity.USER)
val hasClient = entityMap.contains(ClientQuotaEntity.CLIENT_ID)
Expand Down Expand Up @@ -409,13 +443,16 @@ class ZkMigrationClient(zkClient: KafkaZkClient) extends MigrationClient with Lo

tryWriteEntityConfig(configType.get, path.get, props, create=true, state) match {
case Some(newStateSecondTry) => newStateSecondTry
case None => throw new RuntimeException(
case None => throw new MigrationClientException(
s"Could not write client quotas for $entity on second attempt when using Create instead of SetData")
}
}
}

override def writeProducerId(nextProducerId: Long, state: ZkMigrationLeadershipState): ZkMigrationLeadershipState = {
override def writeProducerId(
nextProducerId: Long,
state: ZkMigrationLeadershipState
): ZkMigrationLeadershipState = wrapZkException {
val newProducerIdBlockData = ProducerIdBlockZNode.generateProducerIdBlockJson(
new ProducerIdsBlock(-1, nextProducerId, ProducerIdsBlock.PRODUCER_ID_BLOCK_SIZE))

Expand All @@ -424,9 +461,11 @@ class ZkMigrationClient(zkClient: KafkaZkClient) extends MigrationClient with Lo
state.withMigrationZkVersion(migrationZkVersion)
}

override def writeConfigs(resource: ConfigResource,
configs: util.Map[String, String],
state: ZkMigrationLeadershipState): ZkMigrationLeadershipState = {
override def writeConfigs(
resource: ConfigResource,
configs: util.Map[String, String],
state: ZkMigrationLeadershipState
): ZkMigrationLeadershipState = wrapZkException {
val configType = resource.`type`() match {
case ConfigResource.Type.BROKER => Some(ConfigType.Broker)
case ConfigResource.Type.TOPIC => Some(ConfigType.Topic)
Expand All @@ -447,7 +486,7 @@ class ZkMigrationClient(zkClient: KafkaZkClient) extends MigrationClient with Lo

tryWriteEntityConfig(configType.get, configName, props, create=true, state) match {
case Some(newStateSecondTry) => newStateSecondTry
case None => throw new RuntimeException(
case None => throw new MigrationClientException(
s"Could not write ${configType.get} configs on second attempt when using Create instead of SetData.")
}
}
Expand All @@ -456,10 +495,4 @@ class ZkMigrationClient(zkClient: KafkaZkClient) extends MigrationClient with Lo
state
}
}

override def writeMetadataDeltaToZookeeper(delta: MetadataDelta,
image: MetadataImage,
state: ZkMigrationLeadershipState): ZkMigrationLeadershipState = {
state
}
}
2 changes: 1 addition & 1 deletion core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ class ZkMigrationIntegrationTest {
quotas.add(new ClientQuotaAlteration(
new ClientQuotaEntity(Map("ip" -> "8.8.8.8").asJava),
List(new ClientQuotaAlteration.Op("connection_creation_rate", 10.0)).asJava))
admin.alterClientQuotas(quotas)
admin.alterClientQuotas(quotas).all().get(60, TimeUnit.SECONDS)
Comment thread
mumrah marked this conversation as resolved.

val zkClient = clusterInstance.asInstanceOf[ZkClusterInstance].getUnderlying().zkClient
val migrationClient = new ZkMigrationClient(zkClient)
Expand Down
Loading