KAFKA-15860: ControllerRegistration must be written out to the metadata image

metadata/src/main/java/org/apache/kafka/image/ClusterImage.java

@@ -25,6 +25,7 @@
 
 import java.util.Collections;
 import java.util.Map;
+import java.util.Objects;
 
 
 /**
@@ -73,18 +74,28 @@ public void write(ImageWriter writer, ImageWriterOptions options) {
         for (BrokerRegistration broker : brokers.values()) {
             writer.write(broker.toRecord(options));
         }
+        if (!controllers.isEmpty()) {
+            if (!options.metadataVersion().isControllerRegistrationSupported()) {
+                options.handleLoss("controller registration data");
+            } else {
+                for (ControllerRegistration controller : controllers.values()) {
+                    writer.write(controller.toRecord(options));
+                }
+            }
+        }
     }
 
     @Override
     public int hashCode() {
-        return brokers.hashCode();
+        return Objects.hash(brokers, controllers);
     }
 
     @Override
     public boolean equals(Object o) {
         if (!(o instanceof ClusterImage)) return false;
         ClusterImage other = (ClusterImage) o;
-        return brokers.equals(other.brokers);
+        return brokers.equals(other.brokers) &&
+            controllers.equals(other.controllers);
     }
 
     @Override

metadata/src/main/java/org/apache/kafka/image/writer/ImageWriterOptions.java

@@ -43,27 +43,17 @@ public Builder(MetadataImage image) {
         }
 
         public Builder setMetadataVersion(MetadataVersion metadataVersion) {
-            setRequestedMetadataVersion(metadataVersion);
+            this.requestedMetadataVersion = metadataVersion;
             if (metadataVersion.isLessThan(MetadataVersion.MINIMUM_BOOTSTRAP_VERSION)) {
                 // When writing an image, all versions less than 3.3-IV0 are treated as 3.0-IV1.
                 // This is because those versions don't support FeatureLevelRecord.
-                setRawMetadataVersion(MetadataVersion.MINIMUM_KRAFT_VERSION);
+                this.metadataVersion = MetadataVersion.MINIMUM_KRAFT_VERSION;
             } else {
-                setRawMetadataVersion(metadataVersion);
+                this.metadataVersion = metadataVersion;
             }
             return this;
         }
 
-        // Visible for testing
-        public Builder setRawMetadataVersion(MetadataVersion metadataVersion) {
-            this.metadataVersion = metadataVersion;
-            return this;
-        }
-
-        public void setRequestedMetadataVersion(MetadataVersion orgMetadataVersion) {
-            this.requestedMetadataVersion = orgMetadataVersion;
-        }
-
         public MetadataVersion metadataVersion() {
             return metadataVersion;
         }

metadata/src/test/java/org/apache/kafka/image/ClusterImageTest.java

@@ -35,6 +35,7 @@
 import org.apache.kafka.metadata.RecordTestUtils;
 import org.apache.kafka.metadata.VersionRange;
 import org.apache.kafka.server.common.ApiMessageAndVersion;
+import org.apache.kafka.server.common.MetadataVersion;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 
@@ -45,6 +46,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static org.apache.kafka.common.metadata.MetadataRecordType.FENCE_BROKER_RECORD;
 import static org.apache.kafka.common.metadata.MetadataRecordType.UNFENCE_BROKER_RECORD;
@@ -115,11 +117,11 @@ public class ClusterImageTest {
             (short) 0));
 
         ControllerEndpointCollection endpointsFor1001 = new ControllerEndpointCollection();
-        new ControllerEndpointCollection().add(new ControllerEndpoint().
-                setHost("localhost").
-                setName("PLAINTEXT").
-                setPort(19093).
-                setSecurityProtocol(SecurityProtocol.PLAINTEXT.id));
+        endpointsFor1001.add(new ControllerEndpoint().
+            setHost("localhost").
+            setName("PLAINTEXT").
+            setPort(19093).
+            setSecurityProtocol(SecurityProtocol.PLAINTEXT.id));
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new RegisterControllerRecord().
             setControllerId(1001).
             setIncarnationId(Uuid.fromString("FdEHF-IqScKfYyjZ1CjfNQ")).
@@ -205,4 +207,22 @@ private static List<ApiMessageAndVersion> getImageRecords(ClusterImage image) {
         image.write(writer, new ImageWriterOptions.Builder().build());
         return writer.records();
     }
+
+    @Test
+    public void testHandleLossOfControllerRegistrations() {
+        ClusterImage testImage = new ClusterImage(Collections.emptyMap(),
+            Collections.singletonMap(1000, new ControllerRegistration.Builder().
+                setId(1000).
+                setIncarnationId(Uuid.fromString("9ABu6HEgRuS-hjHLgC4cHw")).
+                setListeners(Collections.singletonMap("PLAINTEXT",
+                    new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 19092))).
+                setSupportedFeatures(Collections.emptyMap()).build()));
+        RecordListWriter writer = new RecordListWriter();
+        final AtomicReference<String> lossString = new AtomicReference<>("");
+        testImage.write(writer, new ImageWriterOptions.Builder().
+            setMetadataVersion(MetadataVersion.IBP_3_6_IV2).
+            setLossHandler(loss -> lossString.compareAndSet("controller registration data", loss.loss())).
+                build());
+        assertEquals("", lossString.get());
+    }
 }