KAFKA-15514: Metadata records Replicas->Assignment

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

@@ -465,6 +465,7 @@ abstract class AbstractControllerBrokerRequestBatch(config: KafkaConfig,
       metadataInstance.partitionLeadershipInfo(partition) match {
         case Some(LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)) =>
           val replicas = metadataInstance.partitionReplicaAssignment(partition)
+          // TODO KAFKA-15362 offline dirs also need to be considered to determine offlineReplicas
           val offlineReplicas = replicas.filter(!metadataInstance.isReplicaOnline(_, partition))
           val updatedLeaderAndIsr =
             if (beingDeleted) LeaderAndIsr.duringDelete(leaderAndIsr.isr)

core/src/main/scala/kafka/migration/MigrationControllerChannelContext.scala

@@ -33,7 +33,7 @@ object MigrationControllerChannelContext {
   def partitionReplicaAssignment(image: MetadataImage, tp: TopicPartition): collection.Seq[Int] = {
     image.topics().topicsByName().asScala.get(tp.topic()) match {
       case Some(topic) => topic.partitions().asScala.get(tp.partition()) match {
-        case Some(partition) => partition.replicas.toSeq
+        case Some(partition) => partition.replicaBrokerIds().toSeq
         case None => collection.Seq.empty
       }
       case None => collection.Seq.empty

core/src/main/scala/kafka/migration/MigrationPropagator.scala

@@ -143,13 +143,13 @@ class MigrationPropagator(
     if (changedZkBrokers.nonEmpty) {
       // For new the brokers, check if there are partition assignments and add LISR appropriately.
       materializePartitions(image.topics()).asScala.foreach { case (tp, partitionRegistration) =>
-        val replicas = partitionRegistration.replicas.toSet
+        val replicas = partitionRegistration.replicaBrokerIds()
         val leaderIsrAndControllerEpochOpt = MigrationControllerChannelContext.partitionLeadershipInfo(image, tp)
-        val newBrokersWithReplicas = replicas.intersect(changedZkBrokers)
+        val newBrokersWithReplicas = replicas.toSet.intersect(changedZkBrokers)
         if (newBrokersWithReplicas.nonEmpty) {
           leaderIsrAndControllerEpochOpt match {
             case Some(leaderIsrAndControllerEpoch) =>
-              val replicaAssignment = ReplicaAssignment(partitionRegistration.replicas,
+              val replicaAssignment = ReplicaAssignment(replicas,
                 partitionRegistration.addingReplicas, partitionRegistration.removingReplicas)
               requestBatch.addLeaderAndIsrRequestForBrokers(newBrokersWithReplicas.toSeq, tp,
                 leaderIsrAndControllerEpoch, replicaAssignment, isNew = true)
@@ -169,14 +169,15 @@ class MigrationPropagator(
       val deletedTopic = delta.image().topics().getTopic(deletedTopicId)
       deletedTopic.partitions().asScala.foreach { case (partition, partitionRegistration) =>
         val tp = new TopicPartition(deletedTopic.name(), partition)
-        val offlineReplicas = partitionRegistration.replicas.filter {
-          MigrationControllerChannelContext.isReplicaOnline(image, _, partitionRegistration.replicas.toSet)
+        val replicas = partitionRegistration.replicaBrokerIds()
+        val offlineReplicas = replicas.filter {
+          MigrationControllerChannelContext.isReplicaOnline(image, _, replicas.toSet)
         }
         val deletedLeaderAndIsr = LeaderAndIsr.duringDelete(partitionRegistration.isr.toList)
-        requestBatch.addStopReplicaRequestForBrokers(partitionRegistration.replicas, tp, deletePartition = true)
+        requestBatch.addStopReplicaRequestForBrokers(replicas, tp, deletePartition = true)
         requestBatch.addUpdateMetadataRequestForBrokers(
           oldZkBrokers.toSeq, zkControllerEpoch, tp, deletedLeaderAndIsr.leader, deletedLeaderAndIsr.leaderEpoch,
-          deletedLeaderAndIsr.partitionEpoch, deletedLeaderAndIsr.isr, partitionRegistration.replicas, offlineReplicas)
+          deletedLeaderAndIsr.partitionEpoch, deletedLeaderAndIsr.isr, replicas, offlineReplicas)
       }
     }
 
@@ -190,7 +191,8 @@ class MigrationPropagator(
         val leaderIsrAndControllerEpochOpt = MigrationControllerChannelContext.partitionLeadershipInfo(image, tp)
         leaderIsrAndControllerEpochOpt match {
           case Some(leaderIsrAndControllerEpoch) =>
-            val replicaAssignment = ReplicaAssignment(partitionRegistration.replicas,
+            val replicas = partitionRegistration.replicaBrokerIds()
+            val replicaAssignment = ReplicaAssignment(replicas,
               partitionRegistration.addingReplicas, partitionRegistration.removingReplicas)
             requestBatch.addLeaderAndIsrRequestForBrokers(replicaAssignment.replicas, tp,
               leaderIsrAndControllerEpoch, replicaAssignment, isNew = true)
@@ -200,9 +202,9 @@ class MigrationPropagator(
         // Check for removed replicas.
         val oldReplicas =
           Option(delta.image().topics().getPartition(topicDelta.id(), tp.partition()))
-            .map(_.replicas.toSet)
+            .map(_.replicaBrokerIds().toSet)
             .getOrElse(Set.empty)
-        val newReplicas = partitionRegistration.replicas.toSet
+        val newReplicas = partitionRegistration.replicaBrokerIds().toSet
         val removedReplicas = oldReplicas -- newReplicas
         if (removedReplicas.nonEmpty) {
           requestBatch.addStopReplicaRequestForBrokers(removedReplicas.toSeq, tp, deletePartition = false)
@@ -233,7 +235,7 @@ class MigrationPropagator(
       val leaderIsrAndControllerEpochOpt = MigrationControllerChannelContext.partitionLeadershipInfo(image, tp)
       leaderIsrAndControllerEpochOpt match {
         case Some(leaderIsrAndControllerEpoch) =>
-          val replicaAssignment = ReplicaAssignment(partitionRegistration.replicas,
+          val replicaAssignment = ReplicaAssignment(partitionRegistration.replicaBrokerIds(),
             partitionRegistration.addingReplicas, partitionRegistration.removingReplicas)
           requestBatch.addLeaderAndIsrRequestForBrokers(replicaAssignment.replicas, tp,
             leaderIsrAndControllerEpoch, replicaAssignment, isNew = true)

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

@@ -285,7 +285,8 @@ class ControllerServer(
             config.passwordEncoderIterations)
           case None => PasswordEncoder.noop()
         }
-        val migrationClient = ZkMigrationClient(zkClient, zkConfigEncoder)
+        val metadataVersion = controller.asInstanceOf[QuorumController].featureControl().metadataVersion()
+        val migrationClient = ZkMigrationClient(zkClient, zkConfigEncoder, metadataVersion)
         val propagator: LegacyPropagator = new MigrationPropagator(config.nodeId, config)
         val migrationDriver = KRaftMigrationDriver.newBuilder()
           .setNodeId(config.nodeId)

core/src/main/scala/kafka/server/metadata/BrokerMetadataPublisher.scala

@@ -79,9 +79,10 @@ object BrokerMetadataPublisher extends Logging {
       val partitionId = log.topicPartition.partition()
       Option(newTopicsImage.getPartition(topicId, partitionId)) match {
         case Some(partition) =>
-          if (!partition.replicas.contains(brokerId)) {
-            info(s"Found stray log dir $log: the current replica assignment ${partition.replicas} " +
-              s"does not contain the local brokerId $brokerId.")
+          val replicas = partition.replicaBrokerIds()
+          if (!replicas.contains(brokerId)) {
+            info(s"Found stray log dir $log: the current replica assignment " +
+              s"${replicas.mkString("(", ", ", ")")} does not contain the local brokerId $brokerId.")
             Some(log.topicPartition)
           } else {
             None

core/src/main/scala/kafka/server/metadata/KRaftMetadataCache.scala

@@ -63,6 +63,7 @@ class KRaftMetadataCache(val brokerId: Int) extends MetadataCache with Logging w
                                        brokers: Array[Int],
                                        listenerName: ListenerName,
                                        filterUnavailableEndpoints: Boolean): java.util.List[Integer] = {
+    // TODO KAFKA-15362
     if (!filterUnavailableEndpoints) {
       Replicas.toList(brokers)
     } else {
@@ -90,7 +91,7 @@ class KRaftMetadataCache(val brokerId: Int) extends MetadataCache with Logging w
       case Some(topic) => Some(topic.partitions().entrySet().asScala.map { entry =>
         val partitionId = entry.getKey
         val partition = entry.getValue
-        val filteredReplicas = maybeFilterAliveReplicas(image, partition.replicas,
+        val filteredReplicas = maybeFilterAliveReplicas(image, partition.replicaBrokerIds(),
           listenerName, errorUnavailableEndpoints)
         val filteredIsr = maybeFilterAliveReplicas(image, partition.isr, listenerName,
           errorUnavailableEndpoints)
@@ -143,11 +144,9 @@ class KRaftMetadataCache(val brokerId: Int) extends MetadataCache with Logging w
   private def getOfflineReplicas(image: MetadataImage,
                                  partition: PartitionRegistration,
                                  listenerName: ListenerName): util.List[Integer] = {
-    // TODO: in order to really implement this correctly, we would need JBOD support.
-    // That would require us to track which replicas were offline on a per-replica basis.
-    // See KAFKA-13005.
+    // TODO KAFKA-15362 (consider offline log directories for each replica)
     val offlineReplicas = new util.ArrayList[Integer](0)
-    for (brokerId <- partition.replicas) {
+    for (brokerId <- partition.replicaBrokerIds()) {
       Option(image.cluster().broker(brokerId)) match {
         case None => offlineReplicas.add(brokerId)
         case Some(broker) => if (broker.fenced() || !broker.listeners().containsKey(listenerName.value())) {
@@ -240,7 +239,7 @@ class KRaftMetadataCache(val brokerId: Int) extends MetadataCache with Logging w
         setLeaderEpoch(partition.leaderEpoch).
         setIsr(Replicas.toList(partition.isr)).
         setZkVersion(partition.partitionEpoch).
-        setReplicas(Replicas.toList(partition.replicas))))
+        setReplicas(Replicas.brokerIdsList(partition.replicas))))
   }
 
   override def numPartitions(topicName: String): Option[Int] = {
@@ -279,7 +278,7 @@ class KRaftMetadataCache(val brokerId: Int) extends MetadataCache with Logging w
     val result = new mutable.HashMap[Int, Node]()
     Option(image.topics().getTopic(tp.topic())).foreach { topic =>
       topic.partitions().values().forEach { partition =>
-        partition.replicas.foreach { replicaId =>
+        partition.replicaBrokerIds().foreach { replicaId =>
           result.put(replicaId, Option(image.cluster().broker(replicaId)) match {
             case None => Node.noNode()
             case Some(broker) if broker.fenced() => Node.noNode()
@@ -345,7 +344,7 @@ class KRaftMetadataCache(val brokerId: Int) extends MetadataCache with Logging w
         partitionInfos.add(new PartitionInfo(topic.name(),
           partitionId,
           node(partition.leader),
-          partition.replicas.map(replica => node(replica)),
+          partition.replicaBrokerIds().map(replica => node(replica)),
           partition.isr.map(replica => node(replica)),
           getOfflineReplicas(image, partition, listenerName).asScala.
             map(replica => node(replica)).toArray))

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

@@ -28,12 +28,11 @@ import org.apache.kafka.common.metadata._
 import org.apache.kafka.common.resource.ResourcePattern
 import org.apache.kafka.common.security.scram.ScramCredential
 import org.apache.kafka.common.{TopicIdPartition, Uuid}
-import org.apache.kafka.metadata.DelegationTokenData
-import org.apache.kafka.metadata.PartitionRegistration
+import org.apache.kafka.metadata.{DelegationTokenData, PartitionRegistration, Replicas}
 import org.apache.kafka.metadata.migration.ConfigMigrationClient.ClientQuotaVisitor
 import org.apache.kafka.metadata.migration.TopicMigrationClient.{TopicVisitor, TopicVisitorInterest}
 import org.apache.kafka.metadata.migration._
-import org.apache.kafka.server.common.{ApiMessageAndVersion, ProducerIdsBlock}
+import org.apache.kafka.server.common.{ApiMessageAndVersion, MetadataVersion, ProducerIdsBlock}
 import org.apache.zookeeper.KeeperException
 import org.apache.zookeeper.KeeperException.{AuthFailedException, NoAuthException, SessionClosedRequireAuthException}
 
@@ -49,13 +48,14 @@ object ZkMigrationClient {
 
   def apply(
     zkClient: KafkaZkClient,
-    zkConfigEncoder: PasswordEncoder
+    zkConfigEncoder: PasswordEncoder,
+    metadataVersion: MetadataVersion
   ): ZkMigrationClient = {
-    val topicClient = new ZkTopicMigrationClient(zkClient)
+    val topicClient = new ZkTopicMigrationClient(zkClient, metadataVersion)
     val configClient = new ZkConfigMigrationClient(zkClient, zkConfigEncoder)
     val aclClient = new ZkAclMigrationClient(zkClient)
     val delegationTokenClient = new ZkDelegationTokenMigrationClient(zkClient)
-    new ZkMigrationClient(zkClient, topicClient, configClient, aclClient, delegationTokenClient)
+    new ZkMigrationClient(zkClient, topicClient, configClient, aclClient, delegationTokenClient, metadataVersion)
   }
 
   /**
@@ -99,7 +99,8 @@ class ZkMigrationClient(
   topicClient: TopicMigrationClient,
   configClient: ConfigMigrationClient,
   aclClient: AclMigrationClient,
-  delegationTokenClient: DelegationTokenMigrationClient
+  delegationTokenClient: DelegationTokenMigrationClient,
+  metadataVersion: MetadataVersion
 ) extends MigrationClient with Logging {
 
   override def getOrCreateMigrationRecoveryState(
@@ -175,15 +176,19 @@ class ZkMigrationClient(
           val record = new PartitionRecord()
             .setTopicId(topicIdPartition.topicId())
             .setPartitionId(topicIdPartition.partition())
-            .setReplicas(partitionRegistration.replicas.map(Integer.valueOf).toList.asJava)
             .setAddingReplicas(partitionRegistration.addingReplicas.map(Integer.valueOf).toList.asJava)
             .setRemovingReplicas(partitionRegistration.removingReplicas.map(Integer.valueOf).toList.asJava)
             .setIsr(partitionRegistration.isr.map(Integer.valueOf).toList.asJava)
             .setLeader(partitionRegistration.leader)
             .setLeaderEpoch(partitionRegistration.leaderEpoch)
             .setPartitionEpoch(partitionRegistration.partitionEpoch)
             .setLeaderRecoveryState(partitionRegistration.leaderRecoveryState.value())
-          partitionRegistration.replicas.foreach(brokerIdConsumer.accept(_))
+          if (metadataVersion.isDirectoryAssignmentSupported) {
+            record.setAssignment(Replicas.toPartitionRecordReplicaAssignment(partitionRegistration.replicas))
+          } else {
+            record.setReplicas(partitionRegistration.replicaBrokerIds().map(Integer.valueOf).toList.asJava)
+          }
+          partitionRegistration.replicaBrokerIds().foreach(brokerIdConsumer.accept(_))
           partitionRegistration.addingReplicas.foreach(brokerIdConsumer.accept(_))
           topicBatch.add(new ApiMessageAndVersion(record, 0.toShort))
         }

core/src/main/scala/kafka/zk/migration/ZkTopicMigrationClient.scala

@@ -29,7 +29,8 @@ import org.apache.kafka.common.metadata.PartitionRecord
 import org.apache.kafka.common.{TopicIdPartition, TopicPartition, Uuid}
 import org.apache.kafka.metadata.migration.TopicMigrationClient.TopicVisitorInterest
 import org.apache.kafka.metadata.migration.{MigrationClientException, TopicMigrationClient, ZkMigrationLeadershipState}
-import org.apache.kafka.metadata.{LeaderRecoveryState, PartitionRegistration}
+import org.apache.kafka.metadata.{LeaderRecoveryState, PartitionRegistration, Replicas}
+import org.apache.kafka.server.common.MetadataVersion
 import org.apache.zookeeper.CreateMode
 import org.apache.zookeeper.KeeperException.Code
 
@@ -39,7 +40,10 @@ import scala.collection.mutable.ArrayBuffer
 import scala.jdk.CollectionConverters._
 
 
-class ZkTopicMigrationClient(zkClient: KafkaZkClient) extends TopicMigrationClient with Logging {
+class ZkTopicMigrationClient(
+    zkClient: KafkaZkClient,
+    metadataVersion: MetadataVersion
+  ) extends TopicMigrationClient with Logging {
   override def iterateTopics(
     interests: util.EnumSet[TopicVisitorInterest],
     visitor: TopicMigrationClient.TopicVisitor,
@@ -60,13 +64,18 @@ class ZkTopicMigrationClient(zkClient: KafkaZkClient) extends TopicMigrationClie
         val partitions = partitionAssignments.keys.toSeq
         val leaderIsrAndControllerEpochs = zkClient.getTopicPartitionStates(partitions)
         partitionAssignments.foreach { case (topicPartition, replicaAssignment) =>
-          val replicaList = replicaAssignment.replicas.map(Integer.valueOf).asJava
           val record = new PartitionRecord()
             .setTopicId(topicIdOpt.get)
             .setPartitionId(topicPartition.partition)
-            .setReplicas(replicaList)
             .setAddingReplicas(replicaAssignment.addingReplicas.map(Integer.valueOf).asJava)
             .setRemovingReplicas(replicaAssignment.removingReplicas.map(Integer.valueOf).asJava)
+          val replicaList = replicaAssignment.replicas.map(Integer.valueOf).asJava
+          if (metadataVersion.isDirectoryAssignmentSupported) {
+            record.setAssignment(Replicas.toPartitionRecordReplicaAssignment(
+              Replicas.withUnknownDirs(replicaList)))
+          } else {
+            record.setReplicas(replicaList)
+          }
           leaderIsrAndControllerEpochs.get(topicPartition) match {
             case Some(leaderIsrAndEpoch) =>
               record
@@ -102,7 +111,7 @@ class ZkTopicMigrationClient(zkClient: KafkaZkClient) extends TopicMigrationClie
 
     val assignments = partitions.asScala.map { case (partitionId, partition) =>
       new TopicPartition(topicName, partitionId) ->
-        ReplicaAssignment(partition.replicas, partition.addingReplicas, partition.removingReplicas)
+        ReplicaAssignment(partition.replicaBrokerIds(), partition.addingReplicas, partition.removingReplicas)
     }
 
     val createTopicZNode = {
@@ -177,7 +186,7 @@ class ZkTopicMigrationClient(zkClient: KafkaZkClient) extends TopicMigrationClie
   ): ZkMigrationLeadershipState = wrapZkException {
     val assignments = partitions.asScala.map { case (partitionId, partition) =>
       new TopicPartition(topicName, partitionId) ->
-        ReplicaAssignment(partition.replicas, partition.addingReplicas, partition.removingReplicas)
+        ReplicaAssignment(partition.replicaBrokerIds(), partition.addingReplicas, partition.removingReplicas)
     }
     val request = SetDataRequest(
       TopicZNode.path(topicName),

core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala

@@ -128,7 +128,7 @@ class ZkMigrationIntegrationTest {
 
     val underlying = clusterInstance.asInstanceOf[ZkClusterInstance].getUnderlying()
     val zkClient = underlying.zkClient
-    val migrationClient = ZkMigrationClient(zkClient, PasswordEncoder.noop())
+    val migrationClient = ZkMigrationClient(zkClient, PasswordEncoder.noop(), MetadataVersion.latest())
     val verifier = new MetadataDeltaVerifier()
     migrationClient.readAllMetadata(batch => verifier.accept(batch), _ => { })
     verifier.verify { image =>
@@ -238,7 +238,7 @@ class ZkMigrationIntegrationTest {
       case None => PasswordEncoder.noop()
     }
 
-    val migrationClient = ZkMigrationClient(zkClient, zkConfigEncoder)
+    val migrationClient = ZkMigrationClient(zkClient, zkConfigEncoder, MetadataVersion.latest())
     var migrationState = migrationClient.getOrCreateMigrationRecoveryState(ZkMigrationLeadershipState.EMPTY)
     migrationState = migrationState.withNewKRaftController(3000, 42)
     migrationState = migrationClient.claimControllerLeadership(migrationState)

core/src/test/scala/unit/kafka/server/ReplicaManagerConcurrencyTest.scala

@@ -437,7 +437,7 @@ class ReplicaManagerConcurrencyTest {
       delta.replay(new PartitionRecord()
         .setTopicId(topic.topicId)
         .setPartitionId(partitionId)
-        .setReplicas(toList(registration.replicas))
+        .setReplicas(toList(registration.replicaBrokerIds()))
         .setIsr(toList(registration.isr))
         .setLeader(registration.leader)
         .setLeaderEpoch(registration.leaderEpoch)
@@ -465,7 +465,7 @@ class ReplicaManagerConcurrencyTest {
     partitionEpoch: Int = 0
   ): PartitionRegistration = {
     new PartitionRegistration.Builder().
-      setReplicas(replicaIds.toArray).
+      setReplicasWithUnknownDirs(replicaIds.toArray).
       setIsr(isr.toArray).
       setLeader(leader).
       setLeaderRecoveryState(leaderRecoveryState).