KAFKA-18225: ClientQuotaCallback#updateClusterMetadata is unsupported by kraft

clients/src/main/java/org/apache/kafka/server/quota/ClientQuotaCallback.java

@@ -23,7 +23,7 @@
 import java.util.Map;
 
 /**
- * Quota callback interface for brokers that enables customization of client quota computation.
+ * Quota callback interface for brokers and controllers that enables customization of client quota computation.
  */
 public interface ClientQuotaCallback extends Configurable {
 
@@ -89,11 +89,11 @@ public interface ClientQuotaCallback extends Configurable {
     boolean quotaResetRequired(ClientQuotaType quotaType);
 
     /**
-     * Metadata update callback that is invoked whenever UpdateMetadata request is received from
-     * the controller. This is useful if quota computation takes partitions into account.
+     * Metadata update callback that is invoked whenever the topic and cluster delta changed. 
+     * This is useful if quota computation takes partitions into account.
      * Topics that are being deleted will not be included in `cluster`.
      *
-     * @param cluster Cluster metadata including partitions and their leaders if known
+     * @param cluster Cluster metadata including topic and cluster
      * @return true if quotas have changed and metric configs may need to be updated
      */
     boolean updateClusterMetadata(Cluster cluster);

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

@@ -507,7 +507,15 @@ class BrokerServer(
           config,
           sharedServer.metadataPublishingFaultHandler,
           "broker",
-          clientQuotaMetadataManager),
+          clientQuotaMetadataManager,
+        ),
+        new DynamicTopicClusterQuotaPublisher(
+          clusterId,
+          config,
+          sharedServer.metadataPublishingFaultHandler,
+          "broker",
+          quotaManagers,
+        ),
         new ScramPublisher(
           config,
           sharedServer.metadataPublishingFaultHandler,

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

@@ -22,7 +22,7 @@ import kafka.raft.KafkaRaftManager
 import kafka.server.QuotaFactory.QuotaManagers
 
 import scala.collection.immutable
-import kafka.server.metadata.{AclPublisher, ClientQuotaMetadataManager, DelegationTokenPublisher, DynamicClientQuotaPublisher, DynamicConfigPublisher, KRaftMetadataCache, KRaftMetadataCachePublisher, ScramPublisher}
+import kafka.server.metadata.{AclPublisher, ClientQuotaMetadataManager, DelegationTokenPublisher, DynamicClientQuotaPublisher, DynamicConfigPublisher, DynamicTopicClusterQuotaPublisher, KRaftMetadataCache, KRaftMetadataCachePublisher, ScramPublisher}
 import kafka.utils.{CoreUtils, Logging}
 import org.apache.kafka.common.message.ApiMessageType.ListenerType
 import org.apache.kafka.common.network.ListenerName
@@ -332,7 +332,17 @@ class ControllerServer(
         config,
         sharedServer.metadataPublishingFaultHandler,
         "controller",
-        clientQuotaMetadataManager))
+        clientQuotaMetadataManager
+      ))
+
+      // Set up the DynamicTopicClusterQuotaPublisher. This will enable quotas for the cluster and topics.
+      metadataPublishers.add(new DynamicTopicClusterQuotaPublisher(
+        clusterId,
+        config,
+        sharedServer.metadataPublishingFaultHandler,
+        "controller",
+        quotaManagers,
+      ))
 
       // Set up the SCRAM publisher.
       metadataPublishers.add(new ScramPublisher(

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

@@ -19,15 +19,20 @@ package kafka.server
 
 import kafka.server.metadata.{ConfigRepository, KRaftMetadataCache}
 import org.apache.kafka.admin.BrokerMetadata
+import org.apache.kafka.common.internals.Topic
 import org.apache.kafka.common.message.{DescribeClientQuotasRequestData, DescribeClientQuotasResponseData, DescribeTopicPartitionsResponseData, DescribeUserScramCredentialsRequestData, DescribeUserScramCredentialsResponseData, MetadataResponseData}
 import org.apache.kafka.common.network.ListenerName
-import org.apache.kafka.common.{Cluster, Node, TopicPartition, Uuid}
-import org.apache.kafka.metadata.LeaderAndIsr
+import org.apache.kafka.common.{Cluster, Node, PartitionInfo, TopicPartition, Uuid}
+import org.apache.kafka.image.MetadataImage
+import org.apache.kafka.metadata.{BrokerRegistration, LeaderAndIsr, PartitionRegistration}
 import org.apache.kafka.server.common.{FinalizedFeatures, KRaftVersion, MetadataVersion}
 
 import java.util
+import java.util.Collections
+import java.util.concurrent.ThreadLocalRandom
 import java.util.function.Supplier
 import scala.collection._
+import scala.jdk.CollectionConverters.CollectionHasAsScala
 
 trait MetadataCache extends ConfigRepository {
   /**
@@ -144,4 +149,95 @@ object MetadataCache {
   ): KRaftMetadataCache = {
     new KRaftMetadataCache(brokerId, kraftVersionSupplier)
   }
+
+  def toCluster(clusterId: String, image: MetadataImage): Cluster = {
+    val brokerToNodes = new util.HashMap[Integer, java.util.List[Node]]
+    image.cluster().brokers()
+      .values().stream()
+      .filter(broker => !broker.fenced())
+      .forEach { broker => brokerToNodes.put(broker.id(), broker.nodes()) }
+
+    def getNodes(id: Int): java.util.List[Node] = brokerToNodes.get(id)
+
+    val partitionInfos = new util.ArrayList[PartitionInfo]
+    val internalTopics = new util.HashSet[String]
+
+    def toArray(replicas: Array[Int]): Array[Node] = {
+      util.Arrays.stream(replicas)
+        .mapToObj(replica => getNodes(replica))
+        .flatMap(replica => replica.stream()).toArray(size => new Array[Node](size))
+    }
+
+    val topicImages = image.topics().topicsByName().values()
+    if (topicImages != null) {
+      topicImages.forEach { topic =>
+        topic.partitions().forEach { (key, value) =>
+          val partitionId = key
+          val partition = value
+          val nodes = getNodes(partition.leader)
+          if (nodes != null) {
+            nodes.forEach(node => {
+              partitionInfos.add(new PartitionInfo(topic.name(),
+                partitionId,
+                node,
+                toArray(partition.replicas),
+                toArray(partition.isr),
+                getOfflineReplicas(image, partition).stream()
+                  .map(replica => getNodes(replica))
+                  .flatMap(replica => replica.stream()).toArray(size => new Array[Node](size))))
+            })
+            if (Topic.isInternal(topic.name())) {
+              internalTopics.add(topic.name())
+            }
+          }
+        }
+      }
+    }
+
+    val controllerNode = getNodes(getRandomAliveBroker(image).getOrElse(-1)) match {
+      case null => Node.noNode()
+      case nodes => nodes.get(0)
+    }
+    // Note: the constructor of Cluster does not allow us to reference unregistered nodes.
+    // So, for example, if partition foo-0 has replicas [1, 2] but broker 2 is not
+    // registered, we pass its replicas as [1, -1]. This doesn't make a lot of sense, but
+    // we are duplicating the behavior of ZkMetadataCache, for now.
+    new Cluster(clusterId, brokerToNodes.values().stream().flatMap(n => n.stream()).collect(util.stream.Collectors.toList()),
+      partitionInfos, Collections.emptySet(), internalTopics, controllerNode)
+  }
+
+  private def getOfflineReplicas(image: MetadataImage,
+                                 partition: PartitionRegistration,
+                                 listenerName: ListenerName = null): util.List[Integer] = {
+    val offlineReplicas = new util.ArrayList[Integer](0)
+    for (brokerId <- partition.replicas) {
+      Option(image.cluster().broker(brokerId)) match {
+        case None => offlineReplicas.add(brokerId)
+        case Some(broker) => if (listenerName == null || isReplicaOffline(partition, listenerName, broker)) {
+          offlineReplicas.add(brokerId)
+        }
+      }
+    }
+    offlineReplicas
+  }
+
+  private def isReplicaOffline(partition: PartitionRegistration, listenerName: ListenerName, broker: BrokerRegistration) =
+    broker.fenced() || !broker.listeners().containsKey(listenerName.value()) || isReplicaInOfflineDir(broker, partition)
+
+  private def isReplicaInOfflineDir(broker: BrokerRegistration, partition: PartitionRegistration): Boolean =
+    !broker.hasOnlineDir(partition.directory(broker.id()))
+
+  private def getRandomAliveBroker(image: MetadataImage): Option[Int] = {
+    val aliveBrokers = getAliveBrokers(image).toList
+    if (aliveBrokers.isEmpty) {
+      None
+    } else {
+      Some(aliveBrokers(ThreadLocalRandom.current().nextInt(aliveBrokers.size)).id)
+    }
+  }
+
+  private def getAliveBrokers(image: MetadataImage): Iterable[BrokerMetadata] = {
+    image.cluster().brokers().values().asScala.filterNot(_.fenced()).
+      map(b => new BrokerMetadata(b.id, b.rack))
+  }
 }

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

@@ -71,6 +71,7 @@ class BrokerMetadataPublisher(
   shareCoordinator: Option[ShareCoordinator],
   var dynamicConfigPublisher: DynamicConfigPublisher,
   dynamicClientQuotaPublisher: DynamicClientQuotaPublisher,
+  dynamicTopicClusterQuotaPublisher: DynamicTopicClusterQuotaPublisher,
   scramPublisher: ScramPublisher,
   delegationTokenPublisher: DelegationTokenPublisher,
   aclPublisher: AclPublisher,
@@ -199,6 +200,9 @@ class BrokerMetadataPublisher(
       // Apply client quotas delta.
       dynamicClientQuotaPublisher.onMetadataUpdate(delta, newImage)
 
+      // Apply topic or cluster quotas delta.
+      dynamicTopicClusterQuotaPublisher.onMetadataUpdate(delta, newImage)
+
       // Apply SCRAM delta.
       scramPublisher.onMetadataUpdate(delta, newImage)
 

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

@@ -0,0 +1,69 @@
+/**
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ **/
+package kafka.server.metadata
+
+import kafka.server.{KafkaConfig, MetadataCache}
+import kafka.server.QuotaFactory.QuotaManagers
+import kafka.utils.Logging
+import org.apache.kafka.image.{MetadataDelta, MetadataImage}
+import org.apache.kafka.image.loader.LoaderManifest
+import org.apache.kafka.server.fault.FaultHandler
+
+/**
+ * Publishing dynamic topic or cluster changes to the client quota manager.
+ * Temporary solution since Cluster objects are immutable and costly to update for every metadata change.
+ * See KAFKA-18239 to trace the issue.
+ */
+class DynamicTopicClusterQuotaPublisher (
+  clusterId: String,
+  conf: KafkaConfig,
+  faultHandler: FaultHandler,
+  nodeType: String,
+  quotaManagers: QuotaManagers
+) extends Logging with org.apache.kafka.image.publisher.MetadataPublisher {
+  logIdent = s"[${name()}] "
+
+  override def name(): String = s"DynamicTopicClusterQuotaPublisher $nodeType id=${conf.nodeId}"
+
+  override def onMetadataUpdate(
+    delta: MetadataDelta,
+    newImage: MetadataImage,
+    manifest: LoaderManifest
+  ): Unit = {
+    onMetadataUpdate(delta, newImage)
+  }
+
+  def onMetadataUpdate(
+    delta: MetadataDelta,
+    newImage: MetadataImage,
+  ): Unit = {
+    val deltaName = s"MetadataDelta up to ${newImage.highestOffsetAndEpoch().offset}"
+    try {
+      quotaManagers.clientQuotaCallback().ifPresent(clientQuotaCallback => {
+        if (delta.topicsDelta() != null || delta.clusterDelta() != null) {
+          val cluster = MetadataCache.toCluster(clusterId, newImage)
+          if (clientQuotaCallback.updateClusterMetadata(cluster)) {
+            quotaManagers.fetch.updateQuotaMetricConfigs()
+            quotaManagers.produce.updateQuotaMetricConfigs()
+            quotaManagers.request.updateQuotaMetricConfigs()
+            quotaManagers.controllerMutation.updateQuotaMetricConfigs()
+          }
+        }
+      })
+    } catch {
+      case t: Throwable => faultHandler.handleFault("Uncaught exception while " +
+        s"publishing dynamic topic or cluster changes from $deltaName", t)
+    }
+  }
+}