apache · mumrah · Apr 26, 2023 · Mar 16, 2023 · Mar 16, 2023 · Mar 16, 2023
diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml
@@ -308,7 +308,7 @@
     <suppress checks="CyclomaticComplexity"
               files="(ClientQuotasImage|KafkaEventQueue|MetadataDelta|QuorumController|ReplicationControlManager).java"/>
     <suppress checks="NPathComplexity"
-              files="(ClientQuotasImage|KafkaEventQueue|ReplicationControlManager|FeatureControlManager|KRaftMigrationDriver).java"/>
+              files="(ClientQuotasImage|KafkaEventQueue|MetadataDelta|ReplicationControlManager|FeatureControlManager|KRaftMigrationDriver).java"/>
     <suppress checks="(NPathComplexity|ClassFanOutComplexity|CyclomaticComplexity|ClassDataAbstractionCoupling|LocalVariableName|MemberName|ParameterName|MethodLength|JavaNCSS|AvoidStarImport)"
             files="metadata[\\/]src[\\/](generated|generated-test)[\\/].+.java$"/>
     <suppress checks="BooleanExpressionComplexity"

diff --git a/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala b/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
@@ -24,13 +24,14 @@ import kafka.test.junit.ClusterTestExtensions
 import kafka.test.junit.ZkClusterInvocationContext.ZkClusterInstance
 import kafka.testkit.{KafkaClusterTestKit, TestKitNodes}
 import kafka.utils.{PasswordEncoder, TestUtils}
-import org.apache.kafka.clients.admin._
+import org.apache.kafka.clients.admin.{Admin, AlterClientQuotasResult, AlterConfigOp, AlterConfigsResult, ConfigEntry, CreateTopicsOptions, NewTopic}
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
 import org.apache.kafka.common.Uuid
 import org.apache.kafka.common.acl.AclOperation.{DESCRIBE, READ, WRITE}
 import org.apache.kafka.common.acl.AclPermissionType.ALLOW
 import org.apache.kafka.common.acl.{AccessControlEntry, AclBinding}
 import org.apache.kafka.common.config.{ConfigResource, TopicConfig}
+import org.apache.kafka.common.errors.TimeoutException
 import org.apache.kafka.common.quota.{ClientQuotaAlteration, ClientQuotaEntity}
 import org.apache.kafka.common.resource.PatternType.{LITERAL, PREFIXED}
 import org.apache.kafka.common.resource.ResourcePattern
@@ -43,13 +44,13 @@ import org.apache.kafka.metadata.authorizer.StandardAcl
 import org.apache.kafka.metadata.migration.ZkMigrationLeadershipState
 import org.apache.kafka.raft.RaftConfig
 import org.apache.kafka.server.common.{ApiMessageAndVersion, MetadataVersion, ProducerIdsBlock}
-import org.junit.jupiter.api.Assertions.{assertEquals, assertFalse, assertNotNull, assertTrue}
+import org.junit.jupiter.api.Assertions.{assertEquals, assertFalse, assertNotNull, assertThrows, assertTrue}
 import org.junit.jupiter.api.Timeout
 import org.junit.jupiter.api.extension.ExtendWith
 import org.slf4j.LoggerFactory
 
 import java.util
-import java.util.concurrent.TimeUnit
+import java.util.concurrent.{ExecutionException, TimeUnit}
 import java.util.{Properties, UUID}
 import scala.collection.Seq
 import scala.jdk.CollectionConverters._
@@ -120,6 +121,11 @@ class ZkMigrationIntegrationTest {
     }
   }
 
+  /**
+   * Test ZkMigrationClient against a real ZooKeeper-backed Kafka cluster. This test creates a ZK cluster
+   * and modifies data using AdminClient. The ZkMigrationClient is then used to read the metadata from ZK
+   * as would happen during a migration. The generated records are then verified.
+   */
   @ClusterTest(brokers = 3, clusterType = Type.ZK, metadataVersion = MetadataVersion.IBP_3_4_IV0)
   def testMigrate(clusterInstance: ClusterInstance): Unit = {
     val admin = clusterInstance.createAdminClient()
@@ -323,4 +329,23 @@ class ZkMigrationIntegrationTest {
       assertTrue(firstProducerIdBlock.firstProducerId() < producerIdBlock.firstProducerId())
     }
   }
+
+  /**
+   * Start a KRaft cluster with migrations enabled, verify that the controller does not accept metadata changes
+   * through the RPCs
+   */
+  @ClusterTest(clusterType = Type.KRAFT, brokers = 0, controllers = 1, metadataVersion = MetadataVersion.IBP_3_4_IV0,
+    serverProperties = Array(new ClusterConfigProperty(key = "zookeeper.metadata.migration.enable", value = "true")))
+  def testPreMigrationMode(clusterInstance: ClusterInstance): Unit = {
+    val adminClient = clusterInstance.createAdminClient()
+    val result = adminClient.createTopics(
+      java.util.Collections.singleton(new NewTopic("test", 10, 1.asInstanceOf[Short])),
+      new CreateTopicsOptions().timeoutMs(1000)
+    )
+
+    assertEquals(
+      classOf[TimeoutException],
+      assertThrows(classOf[ExecutionException], () => result.all().get(5, TimeUnit.MINUTES)).getCause.getClass
+    )
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/server/BrokerRegistrationRequestTest.scala b/core/src/test/scala/unit/kafka/server/BrokerRegistrationRequestTest.scala
@@ -22,6 +22,7 @@ import kafka.test.ClusterInstance
 import kafka.test.annotation.{ClusterConfigProperty, ClusterTest, Type}
 import kafka.test.junit.ClusterTestExtensions
 import org.apache.kafka.clients.ClientResponse
+import org.apache.kafka.common.errors.TimeoutException
 import org.apache.kafka.common.message.{BrokerRegistrationRequestData, BrokerRegistrationResponseData}
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.network.ListenerName
@@ -35,7 +36,7 @@ import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.extension.ExtendWith
 import org.junit.jupiter.api.{Tag, Timeout}
 
-import java.util.concurrent.{CompletableFuture, TimeUnit, TimeoutException}
+import java.util.concurrent.{CompletableFuture, TimeUnit}
 
 /**
  * This test simulates a broker registering with the KRaft quorum under different configurations.

diff --git a/core/src/test/scala/unit/kafka/server/MetadataCacheTest.scala b/core/src/test/scala/unit/kafka/server/MetadataCacheTest.scala
@@ -32,7 +32,7 @@ import org.apache.kafka.common.requests.UpdateMetadataRequest
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.metadata.{BrokerRegistrationChangeRecord, PartitionRecord, RegisterBrokerRecord, RemoveTopicRecord, TopicRecord}
 import org.apache.kafka.common.metadata.RegisterBrokerRecord.{BrokerEndpoint, BrokerEndpointCollection}
-import org.apache.kafka.image.{ClusterImage, MetadataDelta, MetadataImage, MetadataProvenance}
+import org.apache.kafka.image.{ClusterImage, MetadataDelta, MetadataImage, MetadataProvenance, ZkMigrationStateImage}
 import org.apache.kafka.server.common.MetadataVersion
 
 import org.junit.jupiter.api.Assertions._
@@ -72,7 +72,8 @@ object MetadataCacheTest {
           image.clientQuotas(),
           image.producerIds(),
           image.acls(),
-          image.scram())
+          image.scram(),
+          ZkMigrationStateImage.EMPTY)
         val delta = new MetadataDelta.Builder().setImage(partialImage).build()
 
         def toRecord(broker: UpdateMetadataBroker): RegisterBrokerRecord = {

diff --git a/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
@@ -4364,7 +4364,8 @@ class ReplicaManagerTest {
       ClientQuotasImage.EMPTY,
       ProducerIdsImage.EMPTY,
       AclsImage.EMPTY,
-      ScramImage.EMPTY
+      ScramImage.EMPTY,
+      ZkMigrationStateImage.EMPTY
     )
   }
 

diff --git a/metadata/src/main/java/org/apache/kafka/controller/ControllerMetrics.java b/metadata/src/main/java/org/apache/kafka/controller/ControllerMetrics.java
diff --git a/metadata/src/main/java/org/apache/kafka/controller/FeatureControlManager.java b/metadata/src/main/java/org/apache/kafka/controller/FeatureControlManager.java
@@ -52,6 +52,7 @@ public static class Builder {
         private QuorumFeatures quorumFeatures = null;
         private MetadataVersion metadataVersion = MetadataVersion.latest();
         private MetadataVersion minimumBootstrapVersion = MetadataVersion.MINIMUM_BOOTSTRAP_VERSION;
+        private ZkMigrationBootstrap zkMigrationBootstrap = null;
 
         Builder setLogContext(LogContext logContext) {
             this.logContext = logContext;
@@ -78,18 +79,29 @@ Builder setMinimumBootstrapVersion(MetadataVersion minimumBootstrapVersion) {
             return this;
         }
 
+        Builder setZkMigrationBootstrap(ZkMigrationBootstrap zkMigrationBootstrap) {
+            this.zkMigrationBootstrap = zkMigrationBootstrap;
+            return this;
+        }
+
         public FeatureControlManager build() {
             if (logContext == null) logContext = new LogContext();
             if (snapshotRegistry == null) snapshotRegistry = new SnapshotRegistry(logContext);
             if (quorumFeatures == null) {
                 quorumFeatures = new QuorumFeatures(0, new ApiVersions(), QuorumFeatures.defaultFeatureMap(),
                         Collections.emptyList());
             }
-            return new FeatureControlManager(logContext,
+            if (zkMigrationBootstrap == null) {
+                throw new IllegalStateException("ZkMigrationBootstrap cannot be null");
+            }
+            return new FeatureControlManager(
+                logContext,
                 quorumFeatures,
                 snapshotRegistry,
                 metadataVersion,
-                minimumBootstrapVersion);
+                minimumBootstrapVersion,
+                zkMigrationBootstrap
+            );
         }
     }
 
@@ -115,18 +127,22 @@ public FeatureControlManager build() {
      */
     private final MetadataVersion minimumBootstrapVersion;
 
+    private final ZkMigrationBootstrap zkMigrationBootstrap;
+
     private FeatureControlManager(
         LogContext logContext,
         QuorumFeatures quorumFeatures,
         SnapshotRegistry snapshotRegistry,
         MetadataVersion metadataVersion,
-        MetadataVersion minimumBootstrapVersion
+        MetadataVersion minimumBootstrapVersion,
+        ZkMigrationBootstrap zkMigrationBootstrap
     ) {
         this.log = logContext.logger(FeatureControlManager.class);
         this.quorumFeatures = quorumFeatures;
         this.finalizedVersions = new TimelineHashMap<>(snapshotRegistry, 0);
         this.metadataVersion = new TimelineObject<>(snapshotRegistry, metadataVersion);
         this.minimumBootstrapVersion = minimumBootstrapVersion;
+        this.zkMigrationBootstrap = zkMigrationBootstrap;
     }
 
     ControllerResult<Map<String, ApiError>> updateFeatures(
@@ -266,6 +282,11 @@ private ApiError updateMetadataVersion(
             new FeatureLevelRecord()
                 .setName(MetadataVersion.FEATURE_NAME)
                 .setFeatureLevel(newVersionLevel), FEATURE_LEVEL_RECORD.lowestSupportedVersion()));
+
+        // If we are moving to a version that supports migrations, we need to write the correct state into the log
+        if (newVersion.isMigrationSupported()) {
+            zkMigrationBootstrap.bootstrapInitialMigrationState(newVersion, false, recordConsumer);
+        }
         return ApiError.NONE;
     }
 

diff --git a/metadata/src/main/java/org/apache/kafka/controller/ProducerIdControlManager.java b/metadata/src/main/java/org/apache/kafka/controller/ProducerIdControlManager.java
@@ -69,7 +69,7 @@ void replay(ProducerIdsRecord record) {
         ProducerIdsBlock nextBlock = nextProducerBlock.get();
         if (nextBlock != ProducerIdsBlock.EMPTY && record.nextProducerId() <= nextBlock.firstProducerId()) {
             throw new RuntimeException("Next Producer ID from replayed record (" + record.nextProducerId() + ")" +
-                " is not greater than current next Producer ID (" + nextBlock.firstProducerId() + ")");
+                " is not greater than current next Producer ID in block (" + nextBlock + ")");
         } else {
             nextProducerBlock.set(new ProducerIdsBlock(record.brokerId(), record.nextProducerId(), ProducerIdsBlock.PRODUCER_ID_BLOCK_SIZE));
             brokerEpoch.set(record.brokerEpoch());