KAFKA-16308 [1/N]: Create FeatureVersion interface and add --feature flag and handling to StorageTool

core/src/main/scala/kafka/tools/StorageTool.scala

@@ -28,7 +28,7 @@ import net.sourceforge.argparse4j.inf.Namespace
 import org.apache.kafka.common.Uuid
 import org.apache.kafka.common.utils.Utils
 import org.apache.kafka.metadata.bootstrap.{BootstrapDirectory, BootstrapMetadata}
-import org.apache.kafka.server.common.{ApiMessageAndVersion, MetadataVersion}
+import org.apache.kafka.server.common.{ApiMessageAndVersion, FeatureVersion, MetadataVersion}
 import org.apache.kafka.common.metadata.FeatureLevelRecord
 import org.apache.kafka.common.metadata.UserScramCredentialRecord
 import org.apache.kafka.common.security.scram.internals.ScramMechanism
@@ -37,8 +37,7 @@ import org.apache.kafka.metadata.properties.MetaPropertiesEnsemble.VerificationF
 import org.apache.kafka.metadata.properties.{MetaProperties, MetaPropertiesEnsemble, MetaPropertiesVersion, PropertiesUtils}
 
 import java.util
-import java.util.Base64
-import java.util.Optional
+import java.util.{Base64, Collections, Optional}
 import scala.collection.mutable
 import scala.jdk.CollectionConverters._
 import scala.collection.mutable.ArrayBuffer
@@ -59,24 +58,19 @@ object StorageTool extends Logging {
         case "format" =>
           val directories = configToLogDirectories(config.get)
           val clusterId = namespace.getString("cluster_id")
-          val metadataVersion = getMetadataVersion(namespace,
-            Option(config.get.originals.get(KafkaConfig.InterBrokerProtocolVersionProp)).map(_.toString))
-          if (!metadataVersion.isKRaftSupported) {
-            throw new TerseFailure(s"Must specify a valid KRaft metadata.version of at least ${MetadataVersion.IBP_3_0_IV0}.")
-          }
-          if (!metadataVersion.isProduction) {
-            if (config.get.unstableMetadataVersionsEnabled) {
-              System.out.println(s"WARNING: using pre-production metadata.version $metadataVersion.")
-            } else {
-              throw new TerseFailure(s"The metadata.version $metadataVersion is not ready for production use yet.")
-            }
-          }
           val metaProperties = new MetaProperties.Builder().
             setVersion(MetaPropertiesVersion.V1).
             setClusterId(clusterId).
             setNodeId(config.get.nodeId).
             build()
           val metadataRecords : ArrayBuffer[ApiMessageAndVersion] = ArrayBuffer()
+          val specifiedFeatures: util.List[String] = namespace.getList("feature")
+          val featureNamesAndLevelsMap = featureNamesAndLevels(Option(specifiedFeatures).getOrElse(Collections.emptyList).asScala.toList)
+          val metadataVersion = getMetadataVersion(namespace, featureNamesAndLevelsMap,
+            Option(config.get.originals.get(KafkaConfig.InterBrokerProtocolVersionProp)).map(_.toString))
+          metadataVersionValidation(metadataVersion, config)
+          // Get all other features, validate, and create records for them
+          metadataRecords.appendAll(generateFeatureRecords(metadataVersion, featureNamesAndLevelsMap, FeatureVersion.PRODUCTION_FEATURES.asScala.toList))
           getUserScramCredentialRecords(namespace).foreach(userScramCredentialRecords => {
             if (!metadataVersion.isScramSupported) {
               throw new TerseFailure(s"SCRAM is only supported in metadata.version ${MetadataVersion.IBP_3_5_IV2} or later.")
@@ -85,6 +79,7 @@ object StorageTool extends Logging {
               metadataRecords.append(new ApiMessageAndVersion(record, 0.toShort))
             }
           })
+
           val bootstrapMetadata = buildBootstrapMetadata(metadataVersion, Some(metadataRecords), "format command")
           val ignoreFormatted = namespace.getBoolean("ignore_formatted")
           if (!configToSelfManagedMode(config.get)) {
@@ -108,6 +103,43 @@ object StorageTool extends Logging {
     }
   }
 
+  def metadataVersionValidation(metadataVersion: MetadataVersion, config: Option[KafkaConfig]): Unit = {
+    if (!metadataVersion.isKRaftSupported) {
+      throw new TerseFailure(s"Must specify a valid KRaft metadata.version of at least ${MetadataVersion.IBP_3_0_IV0}.")
+    }
+    if (!metadataVersion.isProduction) {
+      if (config.get.unstableMetadataVersionsEnabled) {
+        System.out.println(s"WARNING: using pre-production metadata.version $metadataVersion.")
+      } else {
+        throw new TerseFailure(s"The metadata.version $metadataVersion is not ready for production use yet.")
+      }
+    }
+  }
+
+  def generateFeatureRecords(metadataVersion: MetadataVersion,
+                             specifiedFeatures: Map[String, FeatureVersion],
+                             allFeatures: List[String]): List[ApiMessageAndVersion] = {
+    // If we are using --version-default, the default is based on the metadata version.
+    val metadataVersionOpt: Optional[MetadataVersion] = if (specifiedFeatures.isEmpty) Optional.of(metadataVersion) else Optional.empty[MetadataVersion]
+
+    val featuresMinusMetadata = allFeatures.filter(!_.equals(MetadataVersion.FEATURE_NAME))
+    val allFeaturesAndVersions = featuresMinusMetadata.map { featureName =>
+        specifiedFeatures.getOrElse(featureName, FeatureVersion.defaultValue(featureName, metadataVersionOpt))
+    }
+
+    try {
+      // In order to validate, we need all feature versions set.
+      allFeaturesAndVersions.map { case feature =>
+        feature.validateVersion(metadataVersion, allFeaturesAndVersions.asJava)
+        new ApiMessageAndVersion(new FeatureLevelRecord().
+          setName(feature.featureName).
+          setFeatureLevel(feature.featureLevel), 0.toShort)
+      }
+    } catch {
+      case e: Throwable => throw new TerseFailure(e.getMessage)
+    }
+  }
+
   def parseArguments(args: Array[String]): Namespace = {
     val parser = ArgumentParsers.
       newArgumentParser("kafka-storage", /* defaultHelp */ true, /* prefixChars */ "-", /* fromFilePrefix */ "@").
@@ -140,6 +172,9 @@ object StorageTool extends Logging {
     formatParser.addArgument("--release-version", "-r").
       action(store()).
       help(s"A KRaft release version to use for the initial metadata.version. The minimum is ${MetadataVersion.IBP_3_0_IV0}, the default is ${MetadataVersion.LATEST_PRODUCTION}")
+    formatParser.addArgument("--feature").
+      help("A feature upgrade we should perform, in feature=level format. For example: `metadata.version=5`.").
+      action(append());
 
     parser.parseArgsOrFail(args)
   }
@@ -155,16 +190,27 @@ object StorageTool extends Logging {
 
   def getMetadataVersion(
     namespace: Namespace,
+    featureNamesAndLevelsMap: Map[String, FeatureVersion],
     defaultVersionString: Option[String]
   ): MetadataVersion = {
     val defaultValue = defaultVersionString match {
       case Some(versionString) => MetadataVersion.fromVersionString(versionString)
       case None => MetadataVersion.LATEST_PRODUCTION
     }
 
-    Option(namespace.getString("release_version"))
-      .map(ver => MetadataVersion.fromVersionString(ver))
-      .getOrElse(defaultValue)
+    val releaseVersionTag = Option(namespace.getString("release_version"))
+    val featureTag = featureNamesAndLevelsMap.get(MetadataVersion.FEATURE_NAME)
+
+    (releaseVersionTag, featureTag) match {
+      case (Some(_), Some(_)) =>
+        throw new IllegalArgumentException("Both --release_version and --feature were set. Only one of the two flags can be set.")
+      case (Some(version), None) =>
+        MetadataVersion.fromVersionString(version)
+      case (None, Some(feature)) =>
+        MetadataVersion.fromFeatureLevel(feature.featureLevel())
+      case (None, None) =>
+        defaultValue
+    }
   }
 
   private def getUserScramCredentialRecord(
@@ -464,4 +510,26 @@ object StorageTool extends Logging {
     }
     0
   }
+
+  private def parseNameAndLevel(input: String): Array[String] = {
+    val equalsIndex = input.indexOf("=")
+    if (equalsIndex < 0)
+      throw new RuntimeException("Can't parse feature=level string " + input + ": equals sign not found.")
+    val name = input.substring(0, equalsIndex).trim
+    val levelString = input.substring(equalsIndex + 1).trim
+    try {
+      levelString.toShort
+    } catch {
+      case _: Throwable =>
+        throw new RuntimeException("Can't parse feature=level string " + input + ": " + "unable to parse " + levelString + " as a short.")
+    }
+    Array[String](name, levelString)
+  }
+
+  def featureNamesAndLevels(features: List[String]): Map[String, FeatureVersion] = {
+    features.map((feature: String) => {
+      val nameAndLevel = parseNameAndLevel(feature)
+      (nameAndLevel(0), FeatureVersion.createFeature(nameAndLevel(0), java.lang.Short.valueOf(nameAndLevel(1))))
+    }).toMap
+  }
 }

core/src/test/scala/unit/kafka/tools/StorageToolTest.scala

@@ -21,23 +21,25 @@ import java.io.{ByteArrayOutputStream, PrintStream}
 import java.nio.charset.StandardCharsets
 import java.nio.file.{Files, Paths}
 import java.util
-import java.util.Properties
+import java.util.{Collections, Properties}
 import org.apache.kafka.common.{DirectoryId, KafkaException}
 import kafka.server.KafkaConfig
 import kafka.utils.Exit
 import kafka.utils.TestUtils
+import net.sourceforge.argparse4j.inf.Namespace
 import org.apache.commons.io.output.NullOutputStream
 import org.apache.kafka.common.utils.Utils
-import org.apache.kafka.server.common.MetadataVersion
-import org.apache.kafka.common.metadata.UserScramCredentialRecord
+import org.apache.kafka.server.common.{ApiMessageAndVersion, FeatureVersion, MetadataVersion, TestFeatureVersion}
+import org.apache.kafka.common.metadata.{FeatureLevelRecord, UserScramCredentialRecord}
 import org.apache.kafka.metadata.properties.{MetaProperties, MetaPropertiesEnsemble, MetaPropertiesVersion, PropertiesUtils}
 import org.junit.jupiter.api.Assertions.{assertEquals, assertFalse, assertThrows, assertTrue}
 import org.junit.jupiter.api.{Test, Timeout}
 import org.junit.jupiter.params.ParameterizedTest
-import org.junit.jupiter.params.provider.ValueSource
+import org.junit.jupiter.params.provider.{EnumSource, ValueSource}
 
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters._
 
 @Timeout(value = 40)
 class StorageToolTest {
@@ -52,6 +54,8 @@ class StorageToolTest {
     properties
   }
 
+  val allFeatures = List(MetadataVersion.FEATURE_NAME, TestFeatureVersion.FEATURE_NAME)
+
   @Test
   def testConfigToLogDirectories(): Unit = {
     val config = new KafkaConfig(newSelfManagedProperties())
@@ -204,26 +208,67 @@ Found problem:
   @Test
   def testDefaultMetadataVersion(): Unit = {
     val namespace = StorageTool.parseArguments(Array("format", "-c", "config.props", "-t", "XcZZOzUqS4yHOjhMQB6JLQ"))
-    val mv = StorageTool.getMetadataVersion(namespace, defaultVersionString = None)
+    val mv = StorageTool.getMetadataVersion(namespace, Map.empty, defaultVersionString = None)
     assertEquals(MetadataVersion.LATEST_PRODUCTION.featureLevel(), mv.featureLevel(),
       "Expected the default metadata.version to be the latest production version")
   }
 
   @Test
   def testConfiguredMetadataVersion(): Unit = {
     val namespace = StorageTool.parseArguments(Array("format", "-c", "config.props", "-t", "XcZZOzUqS4yHOjhMQB6JLQ"))
-    val mv = StorageTool.getMetadataVersion(namespace, defaultVersionString = Some(MetadataVersion.IBP_3_3_IV2.toString))
+    val mv = StorageTool.getMetadataVersion(namespace, Map.empty, defaultVersionString = Some(MetadataVersion.IBP_3_3_IV2.toString))
     assertEquals(MetadataVersion.IBP_3_3_IV2.featureLevel(), mv.featureLevel(),
       "Expected the default metadata.version to be 3.3-IV2")
   }
 
+  @Test
+  def testSettingFeatureAndReleaseVersionFails(): Unit = {
+    val namespace = StorageTool.parseArguments(Array("format", "-c", "config.props", "-t", "XcZZOzUqS4yHOjhMQB6JLQ",
+      "--release-version", "3.0-IV1", "--feature", "metadata.version=4"))
+    assertThrows(classOf[IllegalArgumentException], () => StorageTool.getMetadataVersion(namespace, parseFeatures(namespace), defaultVersionString = None))
+  }
+
+  @Test
+  def testParseFeatures(): Unit = {
+    def parseAddFeatures(strings: String*): Map[String, FeatureVersion] = {
+      var args = mutable.Seq("format", "-c", "config.props", "-t", "XcZZOzUqS4yHOjhMQB6JLQ")
+      args ++= strings
+      val namespace = StorageTool.parseArguments(args.toArray)
+      parseFeatures(namespace)
+    }
+
+    assertThrows(classOf[RuntimeException], () => parseAddFeatures("--feature", "blah"))
+    assertThrows(classOf[RuntimeException], () => parseAddFeatures("--feature", "blah=blah"))
+    assertThrows(classOf[IllegalArgumentException], () => parseAddFeatures("--feature", "blah=12"))
+
+    // Test with no features
+    assertEquals(Map(), parseAddFeatures())
+
+    // Test with one feature
+    val testFeatureLevel = 1
+    val testArgument = TestFeatureVersion.FEATURE_NAME + "=" + testFeatureLevel.toString
+    val expectedMap = Map(TestFeatureVersion.FEATURE_NAME -> FeatureVersion.createFeature(TestFeatureVersion.FEATURE_NAME, testFeatureLevel.toShort))
+    assertEquals(expectedMap, parseAddFeatures("--feature", testArgument))
+
+    // Test with two features
+    val metadataFeatureLevel = 5
+    val metadataArgument = MetadataVersion.FEATURE_NAME + "=" + metadataFeatureLevel.toString
+    val expectedMap2 = expectedMap ++ Map (MetadataVersion.FEATURE_NAME -> FeatureVersion.createFeature(MetadataVersion.FEATURE_NAME, metadataFeatureLevel.toShort))
+    assertEquals(expectedMap2, parseAddFeatures("--feature", testArgument, "--feature", metadataArgument))
+  }
+
+  def parseFeatures(namespace: Namespace): Map[String, FeatureVersion] = {
+    val specifiedFeatures: util.List[String] = namespace.getList("feature")
+    StorageTool.featureNamesAndLevels(Option(specifiedFeatures).getOrElse(Collections.emptyList).asScala.toList)
+  }
+
   @Test
   def testMetadataVersionFlags(): Unit = {
     def parseMetadataVersion(strings: String*): MetadataVersion = {
       var args = mutable.Seq("format", "-c", "config.props", "-t", "XcZZOzUqS4yHOjhMQB6JLQ")
       args ++= strings
       val namespace = StorageTool.parseArguments(args.toArray)
-      StorageTool.getMetadataVersion(namespace, defaultVersionString = None)
+      StorageTool.getMetadataVersion(namespace, Map.empty, defaultVersionString = None)
     }
 
     var mv = parseMetadataVersion("--release-version", "3.0")
@@ -235,6 +280,44 @@ Found problem:
     assertThrows(classOf[IllegalArgumentException], () => parseMetadataVersion("--release-version", "0.0"))
   }
 
+  def generateRecord(featureName: String, level: Short): ApiMessageAndVersion = {
+    new ApiMessageAndVersion(new FeatureLevelRecord().
+      setName(featureName).
+      setFeatureLevel(level), 0.toShort)
+  }
+
+  @ParameterizedTest
+  @EnumSource(classOf[TestFeatureVersion])
+  def testFeatureFlag(testFeatureVersion: TestFeatureVersion): Unit = {
+    val featureLevel = testFeatureVersion.featureLevel
+    val records = StorageTool.generateFeatureRecords(MetadataVersion.LATEST_PRODUCTION,
+      Map(TestFeatureVersion.FEATURE_NAME -> FeatureVersion.createFeature(TestFeatureVersion.FEATURE_NAME, featureLevel)),
+      allFeatures
+    )
+    assertEquals(List(generateRecord(TestFeatureVersion.FEATURE_NAME, featureLevel)), records)
+  }
+
+  @ParameterizedTest
+  @EnumSource(classOf[MetadataVersion])
+  def testVersionDefault(metadataVersion: MetadataVersion): Unit = {
+    val records = StorageTool.generateFeatureRecords(metadataVersion,
+      Map.empty,
+      allFeatures
+    )
+
+    val featureLevel = TestFeatureVersion.metadataVersionMapping(metadataVersion).featureLevel()
+    assertEquals(List(generateRecord(TestFeatureVersion.FEATURE_NAME, featureLevel)), records)
+  }
+
+  @Test
+  def testFeatureDependency(): Unit = {
+    val featureLevel = 1.toShort
+    assertThrows(classOf[TerseFailure], () => StorageTool.generateFeatureRecords(MetadataVersion.IBP_2_8_IV1,
+      Map(TestFeatureVersion.FEATURE_NAME -> FeatureVersion.createFeature(TestFeatureVersion.FEATURE_NAME, featureLevel)),
+      allFeatures
+    ))
+  }
+
   @Test
   def testAddScram():Unit = {
     def parseAddScram(strings: String*): Option[ArrayBuffer[UserScramCredentialRecord]] = {