KAFKA-18225: ClientQuotaCallback#updateClusterMetadata is unsupported by kraft

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

@@ -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

@@ -505,10 +505,13 @@ class BrokerServer(
           dynamicConfigHandlers.toMap,
         "broker"),
         new DynamicClientQuotaPublisher(
+          clusterId,
           config,
           sharedServer.metadataPublishingFaultHandler,
           "broker",
-          clientQuotaMetadataManager),
+          clientQuotaMetadataManager,
+          quotaManagers
+        ),
         new ScramPublisher(
           config,
           sharedServer.metadataPublishingFaultHandler,

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

@@ -329,10 +329,13 @@ class ControllerServer(
       // Set up the client quotas publisher. This will enable controller mutation quotas and any
       // other quotas which are applicable.
       metadataPublishers.add(new DynamicClientQuotaPublisher(
+        clusterId,
         config,
         sharedServer.metadataPublishingFaultHandler,
         "controller",
-        clientQuotaMetadataManager))
+        clientQuotaMetadataManager,
+        quotaManagers
+      ))
 
       // Set up the SCRAM publisher.
       metadataPublishers.add(new ScramPublisher(

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

@@ -18,17 +18,20 @@
 package kafka.server.metadata
 
 import kafka.server.KafkaConfig
+import kafka.server.QuotaFactory.QuotaManagers
 import kafka.utils.Logging
 import org.apache.kafka.image.loader.LoaderManifest
 import org.apache.kafka.image.{MetadataDelta, MetadataImage}
 import org.apache.kafka.server.fault.FaultHandler
 
 
 class DynamicClientQuotaPublisher(
+  clusterId: String,
   conf: KafkaConfig,
   faultHandler: FaultHandler,
   nodeType: String,
   clientQuotaMetadataManager: ClientQuotaMetadataManager,
+  quotaManagers: QuotaManagers
 ) extends Logging with org.apache.kafka.image.publisher.MetadataPublisher {
   logIdent = s"[${name()}] "
 
@@ -48,9 +51,20 @@ class DynamicClientQuotaPublisher(
   ): Unit = {
     val deltaName = s"MetadataDelta up to ${newImage.highestOffsetAndEpoch().offset}"
     try {
-        Option(delta.clientQuotasDelta()).foreach { clientQuotasDelta =>
-          clientQuotaMetadataManager.update(clientQuotasDelta)
+      quotaManagers.clientQuotaCallback().ifPresent(clientQuotaCallback => {
+        if (delta.topicsDelta() != null || delta.clusterDelta() != null) {
+          val cluster = KRaftMetadataCache.toCluster(clusterId, newImage)
+          if (clientQuotaCallback.updateClusterMetadata(cluster)) {
+            quotaManagers.fetch.updateQuotaMetricConfigs()
+            quotaManagers.produce.updateQuotaMetricConfigs()
+            quotaManagers.request.updateQuotaMetricConfigs()
+            quotaManagers.controllerMutation.updateQuotaMetricConfigs()
+          }
         }
+      })
+      Option(delta.clientQuotasDelta()).foreach { clientQuotasDelta =>
+        clientQuotaMetadataManager.update(clientQuotasDelta)
+      }
     } catch {
       case t: Throwable => faultHandler.handleFault("Uncaught exception while " +
         s"publishing dynamic client quota changes from $deltaName", t)

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

@@ -17,6 +17,7 @@
 
 package kafka.server.metadata
 
+import kafka.server.metadata.KRaftMetadataCache.{getOfflineReplicas, getRandomAliveBroker}
 import kafka.server.{CachedControllerId, KRaftCachedControllerId, MetadataCache}
 import kafka.utils.Logging
 import org.apache.kafka.admin.BrokerMetadata
@@ -207,28 +208,7 @@ class KRaftMetadataCache(
       }
     }
   }
-
-  private def getOfflineReplicas(image: MetadataImage,
-                                 partition: PartitionRegistration,
-                                 listenerName: ListenerName): 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 (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()))
-
+
   /**
    * Get the endpoint matching the provided listener if the broker is alive. Note that listeners can
    * be added dynamically, so a broker with a missing listener could be a transient error.
@@ -362,12 +342,7 @@ class KRaftMetadataCache(
     Option(_currentImage.cluster.broker(brokerId)).count(_.inControlledShutdown) == 1
   }
 
-  override def getAliveBrokers(): Iterable[BrokerMetadata] = getAliveBrokers(_currentImage)
-
-  private def getAliveBrokers(image: MetadataImage): Iterable[BrokerMetadata] = {
-    image.cluster().brokers().values().asScala.filterNot(_.fenced()).
-      map(b => new BrokerMetadata(b.id, b.rack))
-  }
+  override def getAliveBrokers(): Iterable[BrokerMetadata] = KRaftMetadataCache.getAliveBrokers(_currentImage)
 
   override def getAliveBrokerNode(brokerId: Int, listenerName: ListenerName): Option[Node] = {
     Option(_currentImage.cluster().broker(brokerId)).filterNot(_.fenced()).
@@ -461,15 +436,6 @@ class KRaftMetadataCache(
     getRandomAliveBroker(_currentImage)
   }
 
-  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)
-    }
-  }
-
   def getAliveBrokerEpoch(brokerId: Int): Option[Long] = {
     Option(_currentImage.cluster().broker(brokerId)).filterNot(_.fenced()).
       map(brokerRegistration => brokerRegistration.epoch())
@@ -563,3 +529,96 @@ class KRaftMetadataCache(
   }
 }
 
+object KRaftMetadataCache {
+
+  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))
+  }
+}