KAFKA-15263 Check KRaftMigrationDriver state in each event

metadata/src/main/java/org/apache/kafka/metadata/migration/KRaftMigrationDriver.java

@@ -297,6 +297,16 @@ private boolean isValidStateChange(MigrationDriverState newState) {
         }
     }
 
+    private boolean checkDriverState(MigrationDriverState expectedState) {
+        if (migrationState.equals(expectedState)) {
+            return true;
+        } else {
+            log.debug("Expected driver state {} but found {}. Not running this event {}.",
+                expectedState, migrationState, this.getClass().getSimpleName());
+            return false;
+        }
+    }
+
     private void transitionTo(MigrationDriverState newState) {
         if (!isValidStateChange(newState)) {
             throw new IllegalStateException(
@@ -520,7 +530,7 @@ class WaitForControllerQuorumEvent extends MigrationEvent {
 
         @Override
         public void run() throws Exception {
-            if (migrationState.equals(MigrationDriverState.WAIT_FOR_CONTROLLER_QUORUM)) {
+            if (checkDriverState(MigrationDriverState.WAIT_FOR_CONTROLLER_QUORUM)) {
                 if (!firstPublish) {
                     log.trace("Waiting until we have received metadata before proceeding with migration");
                     return;
@@ -564,24 +574,19 @@ public void run() throws Exception {
     class WaitForZkBrokersEvent extends MigrationEvent {
         @Override
         public void run() throws Exception {
-            switch (migrationState) {
-                case WAIT_FOR_BROKERS:
-                    if (areZkBrokersReadyForMigration()) {
-                        log.info("Zk brokers are registered and ready for migration");
-                        transitionTo(MigrationDriverState.BECOME_CONTROLLER);
-                    }
-                    break;
-                default:
-                    // Ignore the event as we're not in the appropriate state anymore.
-                    break;
+            if (checkDriverState(MigrationDriverState.WAIT_FOR_BROKERS)) {
+                if (areZkBrokersReadyForMigration()) {
+                    log.info("Zk brokers are registered and ready for migration");
+                    transitionTo(MigrationDriverState.BECOME_CONTROLLER);
+                }
             }
         }
     }
 
     class BecomeZkControllerEvent extends MigrationEvent {
         @Override
         public void run() throws Exception {
-            if (migrationState == MigrationDriverState.BECOME_CONTROLLER) {
+            if (checkDriverState(MigrationDriverState.BECOME_CONTROLLER)) {
                 applyMigrationOperation("Claiming ZK controller leadership", zkMigrationClient::claimControllerLeadership);
                 if (migrationLeadershipState.zkControllerEpochZkVersion() == -1) {
                     log.info("Unable to claim leadership, will retry until we learn of a different KRaft leader");
@@ -599,6 +604,9 @@ public void run() throws Exception {
     class MigrateMetadataEvent extends MigrationEvent {
         @Override
         public void run() throws Exception {
+            if (!checkDriverState(MigrationDriverState.ZK_MIGRATION)) {
+                return;
+            }
             Set<Integer> brokersInMetadata = new HashSet<>();
             log.info("Starting ZK migration");
             MigrationManifest.Builder manifestBuilder = MigrationManifest.newBuilder(time);
@@ -655,7 +663,7 @@ public void run() throws Exception {
     class SyncKRaftMetadataEvent extends MigrationEvent {
         @Override
         public void run() throws Exception {
-            if (migrationState == MigrationDriverState.SYNC_KRAFT_TO_ZK) {
+            if (checkDriverState(MigrationDriverState.SYNC_KRAFT_TO_ZK)) {
                 log.info("Performing a full metadata sync from KRaft to ZK.");
                 Map<String, Integer> dualWriteCounts = new TreeMap<>();
                 zkMetadataWriter.handleSnapshot(image, countingOperationConsumer(
@@ -671,7 +679,7 @@ class SendRPCsToBrokersEvent extends MigrationEvent {
         @Override
         public void run() throws Exception {
             // Ignore sending RPCs to the brokers since we're no longer in the state.
-            if (migrationState == MigrationDriverState.KRAFT_CONTROLLER_TO_BROKER_COMM) {
+            if (checkDriverState(MigrationDriverState.KRAFT_CONTROLLER_TO_BROKER_COMM)) {
                 if (image.highestOffsetAndEpoch().compareTo(migrationLeadershipState.offsetAndEpoch()) >= 0) {
                     log.trace("Sending RPCs to broker before moving to dual-write mode using " +
                             "at offset and epoch {}", image.highestOffsetAndEpoch());

metadata/src/test/java/org/apache/kafka/metadata/migration/KRaftMigrationDriverTest.java

@@ -69,6 +69,7 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -634,4 +635,49 @@ public void testControllerFailover() throws Exception {
             assertEquals(new ConfigResource(ConfigResource.Type.TOPIC, "foo"), configClient.deletedResources.get(0));
         });
     }
+
+    @Test
+    public void testBeginMigrationOnce() throws Exception {
+        AtomicInteger migrationBeginCalls = new AtomicInteger(0);
+        NoOpRecordConsumer recordConsumer = new NoOpRecordConsumer() {
+            @Override
+            public void beginMigration() {
+                migrationBeginCalls.incrementAndGet();
+            }
+        };
+        CountingMetadataPropagator metadataPropagator = new CountingMetadataPropagator();
+        CapturingMigrationClient migrationClient = CapturingMigrationClient.newBuilder().setBrokersInZk(1, 2, 3).build();
+        MockFaultHandler faultHandler = new MockFaultHandler("testTwoMigrateMetadataEvents");
+        KRaftMigrationDriver.Builder builder = defaultTestBuilder()
+            .setZkMigrationClient(migrationClient)
+            .setZkRecordConsumer(recordConsumer)
+            .setPropagator(metadataPropagator)
+            .setFaultHandler(faultHandler);
+        try (KRaftMigrationDriver driver = builder.build()) {
+            MetadataImage image = MetadataImage.EMPTY;
+            MetadataDelta delta = new MetadataDelta(image);
+
+            driver.start();
+            delta.replay(ZkMigrationState.PRE_MIGRATION.toRecord().message());
+            delta.replay(zkBrokerRecord(1));
+            delta.replay(zkBrokerRecord(2));
+            delta.replay(zkBrokerRecord(3));
+            MetadataProvenance provenance = new MetadataProvenance(100, 1, 1);
+            image = delta.apply(provenance);
+
+            driver.onControllerChange(new LeaderAndEpoch(OptionalInt.of(3000), 1));
+
+            // Call onMetadataUpdate twice. The first call will trigger the migration to begin (due to presence of brokers)
+            // Both calls will "wakeup" the driver and cause a PollEvent to be run. Calling these back-to-back effectively
+            // causes two MigrateMetadataEvents to be enqueued. Ensure only one is actually run.
+            driver.onMetadataUpdate(delta, image, new LogDeltaManifest(provenance,
+                new LeaderAndEpoch(OptionalInt.of(3000), 1), 1, 100, 42));
+            driver.onMetadataUpdate(delta, image, new LogDeltaManifest(provenance,
+                new LeaderAndEpoch(OptionalInt.of(3000), 1), 1, 100, 42));
+
+            TestUtils.waitForCondition(() -> driver.migrationState().get(1, TimeUnit.MINUTES).equals(MigrationDriverState.DUAL_WRITE),
+                    "Waiting for KRaftMigrationDriver to enter ZK_MIGRATION state");
+            assertEquals(1, migrationBeginCalls.get());
+        }
+    }
 }