From 914a465286179d62b9662595ddc4cb2bd4123361 Mon Sep 17 00:00:00 2001 From: David Arthur Date: Wed, 10 Feb 2021 12:45:43 -0500 Subject: [PATCH 1/2] Bring over files from kip-500 branch --- .../server/metadata/ClientQuotaCache.scala | 284 +++++++++++ .../metadata/ClientQuotaMetadataManager.scala | 174 +++++++ .../ClientQuotaMetadataManagerTest.scala | 452 ++++++++++++++++++ 3 files changed, 910 insertions(+) create mode 100644 core/src/main/scala/kafka/server/metadata/ClientQuotaCache.scala create mode 100644 core/src/main/scala/kafka/server/metadata/ClientQuotaMetadataManager.scala create mode 100644 core/src/test/scala/unit/kafka/server/metadata/ClientQuotaMetadataManagerTest.scala diff --git a/core/src/main/scala/kafka/server/metadata/ClientQuotaCache.scala b/core/src/main/scala/kafka/server/metadata/ClientQuotaCache.scala new file mode 100644 index 0000000000000..aedab2e0769db --- /dev/null +++ b/core/src/main/scala/kafka/server/metadata/ClientQuotaCache.scala @@ -0,0 +1,284 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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.utils.CoreUtils.{inReadLock, inWriteLock} +import org.apache.kafka.common.errors.{InvalidRequestException, UnsupportedVersionException} +import org.apache.kafka.common.quota.{ClientQuotaEntity, ClientQuotaFilterComponent} + +import java.util.concurrent.locks.ReentrantReadWriteLock +import scala.collection.mutable +import scala.jdk.CollectionConverters._ + + +// A type for the cache index keys +sealed trait CacheIndexKey +case object DefaultUser extends CacheIndexKey +case class SpecificUser(user: String) extends CacheIndexKey +case object DefaultClientId extends CacheIndexKey +case class SpecificClientId(clientId: String) extends CacheIndexKey +case object DefaultIp extends CacheIndexKey +case class SpecificIp(ip: String) extends CacheIndexKey + + +// Different types of matching constraints +sealed trait QuotaMatch +case class ExactMatch(entityName: String) extends QuotaMatch +case object DefaultMatch extends QuotaMatch +case object TypeMatch extends QuotaMatch + + +class ClientQuotaCache { + private type QuotaCacheIndex = mutable.HashMap[CacheIndexKey, mutable.HashSet[QuotaEntity]] + + // A mapping of the quota entities to their quotas, for example: + // { + // (user:alice) -> {consumer_byte_rate: 10000}, + // (user:alice,client:x) -> {consumer_byte_rate: 8000, producer_byte_rate: 8000} + // } + private val quotaCache = new mutable.HashMap[QuotaEntity, mutable.Map[String, Double]] + + // Indexes for the three supported entity types. This is needed for flexible lookups like: "all quotas for user:alice" + // The structure of these indexes is a mapping of a specific entity entry to entities in the cache: + // { + // SpecificUser(alice) -> [(user:alice), (user:alice,client:x)], + // DefaultUser -> [(user:default), (user:default, client:default), (user:default, client:x), ...] + // } + // We need three separate indexes because we also support wildcard lookups on entity type. + private val userEntityIndex = new QuotaCacheIndex + private val clientIdEntityIndex = new QuotaCacheIndex + private val ipEntityIndex = new QuotaCacheIndex + + private val lock = new ReentrantReadWriteLock() + + /** + * Return quota entries for a given filter. These entries are returned from an in-memory cache and may not reflect + * the latest state of the quotas according to the controller. If a filter is given for an unsupported entity type + * or an invalid combination of entity types, this method will throw an exception. + * + * @param filters A collection of quota filters (entity type and a match clause). + * @param strict True if we should only return entities which match all the filter clauses and have no + * additional unmatched parts. + * @return A mapping of quota entities along with their quota values. + */ + def describeClientQuotas(filters: Seq[ClientQuotaFilterComponent], strict: Boolean): + Map[ClientQuotaEntity, Map[String, Double]] = inReadLock(lock) { + describeClientQuotasInternal(filters, strict).map { case (entity, value) => convertEntity(entity) -> value} + } + + // Visible for testing (QuotaEntity is nicer for assertions in test code) + private[metadata] def describeClientQuotasInternal(filters: Seq[ClientQuotaFilterComponent], strict: Boolean): + Map[QuotaEntity, Map[String, Double]] = inReadLock(lock) { + + // Do some preliminary validation of the filter types and convert them to correct QuotaMatch type + val entityFilters = mutable.HashMap.empty[String, QuotaMatch] + filters.foreach { component => + val entityType = component.entityType() + if (entityFilters.contains(entityType)) { + throw new InvalidRequestException(s"Duplicate ${entityType} filter component entity type") + } else if (entityType.isEmpty) { + throw new InvalidRequestException("Unexpected empty filter component entity type") + } else if (!ClientQuotaEntity.isValidEntityType(entityType)) { + throw new UnsupportedVersionException(s"Custom entity type ${entityType} not supported") + } + + // A present "match()" is an exact match on name, an absent "match()" is a match on the default entity, + // and a null "match()" is a match on the entity type + val entityMatch = if (component.`match`() != null && component.`match`().isPresent) { + ExactMatch(component.`match`().get()) + } else if (component.`match`() != null) { + DefaultMatch + } else { + TypeMatch + } + entityFilters.put(entityType, entityMatch) + } + + if (entityFilters.isEmpty) { + return Map.empty + } + + // We do not allow IP filters to be combined with user or client filters + val matchingEntities: Set[QuotaEntity] = if (entityFilters.contains(ClientQuotaEntity.IP)) { + if (entityFilters.size > 1) { + throw new InvalidRequestException("Invalid entity filter component combination, IP filter component should " + + "not be used with user or clientId filter component.") + } + val ipMatch = entityFilters.get(ClientQuotaEntity.IP) + ipMatch.fold(Set.empty[QuotaEntity]) { + case ExactMatch(ip) => ipEntityIndex.getOrElse(SpecificIp(ip), Set.empty).toSet + case DefaultMatch => ipEntityIndex.getOrElse(DefaultIp, Set.empty).toSet + case TypeMatch => ipEntityIndex.values.flatten.toSet + } + } else if (entityFilters.contains(ClientQuotaEntity.USER) || entityFilters.contains(ClientQuotaEntity.CLIENT_ID)) { + // If either are present, check both user and client indexes + val userMatch = entityFilters.get(ClientQuotaEntity.USER) + val userIndexMatches = userMatch.fold(Set.empty[QuotaEntity]) { + case ExactMatch(user) => userEntityIndex.getOrElse(SpecificUser(user), Set.empty).toSet + case DefaultMatch => userEntityIndex.getOrElse(DefaultUser, Set.empty).toSet + case TypeMatch => userEntityIndex.values.flatten.toSet + } + + val clientMatch = entityFilters.get(ClientQuotaEntity.CLIENT_ID) + val clientIndexMatches = clientMatch.fold(Set.empty[QuotaEntity]) { + case ExactMatch(clientId) => clientIdEntityIndex.getOrElse(SpecificClientId(clientId), Set.empty).toSet + case DefaultMatch => clientIdEntityIndex.getOrElse(DefaultClientId, Set.empty).toSet + case TypeMatch => clientIdEntityIndex.values.flatten.toSet + } + + val candidateMatches = if (userMatch.isDefined && clientMatch.isDefined) { + userIndexMatches.intersect(clientIndexMatches) + } else if (userMatch.isDefined) { + userIndexMatches + } else { + clientIndexMatches + } + + if (strict) { + // If in strict mode, we need to remove any matches with unspecified entity types. This only applies to results + // with more than one entity part (i.e., user and clientId) + candidateMatches.filter { quotaEntity => + quotaEntity match { + case ExplicitUserExplicitClientIdEntity(_, _) => userMatch.isDefined && clientMatch.isDefined + case DefaultUserExplicitClientIdEntity(_) => userMatch.isDefined && clientMatch.isDefined + case ExplicitUserDefaultClientIdEntity(_) => userMatch.isDefined && clientMatch.isDefined + case DefaultUserDefaultClientIdEntity => userMatch.isDefined && clientMatch.isDefined + case _ => true + } + } + } else { + candidateMatches + } + } else { + // ClientQuotaEntity.isValidEntityType check above should prevent any unknown entity types + throw new IllegalStateException(s"Unexpected handling of ${entityFilters} after filter validation") + } + + val resultsMap: Map[QuotaEntity, Map[String, Double]] = matchingEntities.map { + quotaEntity => { + quotaCache.get(quotaEntity) match { + case Some(quotas) => quotaEntity -> quotas.toMap + case None => quotaEntity -> Map.empty[String, Double] + } + } + }.toMap + + resultsMap + } + + private def convertEntity(entity: QuotaEntity): ClientQuotaEntity = { + val entityMap = entity match { + case IpEntity(ip) => Map(ClientQuotaEntity.IP -> ip) + case DefaultIpEntity => Map(ClientQuotaEntity.IP -> null) + case UserEntity(user) => Map(ClientQuotaEntity.USER -> user) + case DefaultUserEntity => Map(ClientQuotaEntity.USER -> null) + case ClientIdEntity(clientId) => Map(ClientQuotaEntity.CLIENT_ID -> clientId) + case DefaultClientIdEntity => Map(ClientQuotaEntity.CLIENT_ID -> null) + case ExplicitUserExplicitClientIdEntity(user, clientId) => + Map(ClientQuotaEntity.USER -> user, ClientQuotaEntity.CLIENT_ID -> clientId) + case ExplicitUserDefaultClientIdEntity(user) => + Map(ClientQuotaEntity.USER -> user, ClientQuotaEntity.CLIENT_ID -> null) + case DefaultUserExplicitClientIdEntity(clientId) => + Map(ClientQuotaEntity.USER -> null, ClientQuotaEntity.CLIENT_ID -> clientId) + case DefaultUserDefaultClientIdEntity => + Map(ClientQuotaEntity.USER -> null, ClientQuotaEntity.CLIENT_ID -> null) + } + new ClientQuotaEntity(entityMap.asJava) + } + + // Update the cache indexes + private def updateCacheIndex(quotaEntity: QuotaEntity, + remove: Boolean) + (quotaCacheIndex: QuotaCacheIndex, + key: CacheIndexKey): Unit = { + if (remove) { + val needsCleanup = quotaCacheIndex.get(key) match { + case Some(quotaEntitySet) => + quotaEntitySet.remove(quotaEntity) + quotaEntitySet.isEmpty + case None => false + } + if (needsCleanup) { + quotaCacheIndex.remove(key) + } + } else { + quotaCacheIndex.getOrElseUpdate(key, mutable.HashSet.empty).add(quotaEntity) + } + } + + /** + * Update the quota cache with the given entity and quota key/value. If remove is set, the value is ignore and + * the quota entry is removed for the given key. No validation on quota keys is performed here, it is assumed + * that the caller has already done this. + * + * @param entity A quota entity, either a specific entity or the default entity for the given type(s) + * @param key The quota key + * @param value The quota value + * @param remove True if we should remove the given quota key from the entity's quota cache + */ + def updateQuotaCache(entity: QuotaEntity, key: String, value: Double, remove: Boolean): Unit = inWriteLock(lock) { + val quotaValues = quotaCache.getOrElseUpdate(entity, mutable.HashMap.empty) + val removeFromIndex = if (remove) { + quotaValues.remove(key) + if (quotaValues.isEmpty) { + quotaCache.remove(entity) + true + } else { + false + } + } else { + quotaValues.put(key, value) + false + } + + // Update the appropriate indexes with the entity + val updateCacheIndexPartial: (QuotaCacheIndex, CacheIndexKey) => Unit = updateCacheIndex(entity, removeFromIndex) + entity match { + case UserEntity(user) => + updateCacheIndexPartial(userEntityIndex, SpecificUser(user)) + case DefaultUserEntity => + updateCacheIndexPartial(userEntityIndex, DefaultUser) + + case ClientIdEntity(clientId) => + updateCacheIndexPartial(clientIdEntityIndex, SpecificClientId(clientId)) + case DefaultClientIdEntity => + updateCacheIndexPartial(clientIdEntityIndex, DefaultClientId) + + case ExplicitUserExplicitClientIdEntity(user, clientId) => + updateCacheIndexPartial(userEntityIndex, SpecificUser(user)) + updateCacheIndexPartial(clientIdEntityIndex, SpecificClientId(clientId)) + + case ExplicitUserDefaultClientIdEntity(user) => + updateCacheIndexPartial(userEntityIndex, SpecificUser(user)) + updateCacheIndexPartial(clientIdEntityIndex, DefaultClientId) + + case DefaultUserExplicitClientIdEntity(clientId) => + updateCacheIndexPartial(userEntityIndex, DefaultUser) + updateCacheIndexPartial(clientIdEntityIndex, SpecificClientId(clientId)) + + case DefaultUserDefaultClientIdEntity => + updateCacheIndexPartial(userEntityIndex, DefaultUser) + updateCacheIndexPartial(clientIdEntityIndex, DefaultClientId) + + case IpEntity(ip) => + updateCacheIndexPartial(ipEntityIndex, SpecificIp(ip)) + case DefaultIpEntity => + updateCacheIndexPartial(ipEntityIndex, DefaultIp) + } + } +} diff --git a/core/src/main/scala/kafka/server/metadata/ClientQuotaMetadataManager.scala b/core/src/main/scala/kafka/server/metadata/ClientQuotaMetadataManager.scala new file mode 100644 index 0000000000000..bcbae8ae1d85a --- /dev/null +++ b/core/src/main/scala/kafka/server/metadata/ClientQuotaMetadataManager.scala @@ -0,0 +1,174 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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.network.ConnectionQuotas +import kafka.server.ConfigEntityName +import kafka.server.QuotaFactory.QuotaManagers +import kafka.utils.Logging +import org.apache.kafka.common.config.internals.QuotaConfigs +import org.apache.kafka.common.metadata.QuotaRecord +import org.apache.kafka.common.metrics.Quota +import org.apache.kafka.common.quota.ClientQuotaEntity +import org.apache.kafka.common.utils.Sanitizer + +import java.net.{InetAddress, UnknownHostException} +import scala.collection.mutable + + +// A strict hierarchy of entities that we support +sealed trait QuotaEntity +case class IpEntity(ip: String) extends QuotaEntity +case object DefaultIpEntity extends QuotaEntity +case class UserEntity(user: String) extends QuotaEntity +case object DefaultUserEntity extends QuotaEntity +case class ClientIdEntity(clientId: String) extends QuotaEntity +case object DefaultClientIdEntity extends QuotaEntity +case class ExplicitUserExplicitClientIdEntity(user: String, clientId: String) extends QuotaEntity +case class ExplicitUserDefaultClientIdEntity(user: String) extends QuotaEntity +case class DefaultUserExplicitClientIdEntity(clientId: String) extends QuotaEntity +case object DefaultUserDefaultClientIdEntity extends QuotaEntity + +/** + * Watch for changes to quotas in the metadata log and update quota managers and cache as necessary + */ +class ClientQuotaMetadataManager(private[metadata] val quotaManagers: QuotaManagers, + private[metadata] val connectionQuotas: ConnectionQuotas, + private[metadata] val quotaCache: ClientQuotaCache) extends Logging { + + def handleQuotaRecord(quotaRecord: QuotaRecord): Unit = { + val entityMap = mutable.Map[String, String]() + quotaRecord.entity().forEach { entityData => + entityMap.put(entityData.entityType(), entityData.entityName()) + } + + if (entityMap.contains(ClientQuotaEntity.IP)) { + // In the IP quota manager, None is used for default entity + val ipEntity = Option(entityMap(ClientQuotaEntity.IP)) match { + case Some(ip) => IpEntity(ip) + case None => DefaultIpEntity + } + handleIpQuota(ipEntity, quotaRecord) + } else if (entityMap.contains(ClientQuotaEntity.USER) || entityMap.contains(ClientQuotaEntity.CLIENT_ID)) { + // Need to handle null values for default entity name, so use "getOrElse" combined with "contains" checks + val userVal = entityMap.getOrElse(ClientQuotaEntity.USER, null) + val clientIdVal = entityMap.getOrElse(ClientQuotaEntity.CLIENT_ID, null) + + // In User+Client quota managers, "" is used for default entity, so we need to represent all possible + // combinations of values, defaults, and absent entities + val userClientEntity = if (entityMap.contains(ClientQuotaEntity.USER) && entityMap.contains(ClientQuotaEntity.CLIENT_ID)) { + if (userVal == null && clientIdVal == null) { + DefaultUserDefaultClientIdEntity + } else if (userVal == null) { + DefaultUserExplicitClientIdEntity(clientIdVal) + } else if (clientIdVal == null) { + ExplicitUserDefaultClientIdEntity(userVal) + } else { + ExplicitUserExplicitClientIdEntity(userVal, clientIdVal) + } + } else if (entityMap.contains(ClientQuotaEntity.USER)) { + if (userVal == null) { + DefaultUserEntity + } else { + UserEntity(userVal) + } + } else { + if (clientIdVal == null) { + DefaultClientIdEntity + } else { + ClientIdEntity(clientIdVal) + } + } + handleUserClientQuota( + userClientEntity, + quotaRecord + ) + } else { + warn(s"Ignoring unsupported quota entity ${quotaRecord.entity()}") + } + } + + def handleIpQuota(ipEntity: QuotaEntity, quotaRecord: QuotaRecord): Unit = { + val inetAddress = ipEntity match { + case IpEntity(ip) => + try { + Some(InetAddress.getByName(ip)) + } catch { + case _: UnknownHostException => throw new IllegalArgumentException(s"Unable to resolve address $ip") + } + case DefaultIpEntity => None + case _ => throw new IllegalStateException("Should only handle IP quota entities here") + } + + // The connection quota only understands the connection rate limit + if (quotaRecord.key() != QuotaConfigs.IP_CONNECTION_RATE_OVERRIDE_CONFIG) { + warn(s"Ignoring unexpected quota key ${quotaRecord.key()} for entity $ipEntity") + return + } + + // Update the cache + quotaCache.updateQuotaCache(ipEntity, quotaRecord.key, quotaRecord.value, quotaRecord.remove) + + // Convert the value to an appropriate Option for the quota manager + val newValue = if (quotaRecord.remove()) { + None + } else { + Some(quotaRecord.value).map(_.toInt) + } + connectionQuotas.updateIpConnectionRateQuota(inetAddress, newValue) + } + + def handleUserClientQuota(quotaEntity: QuotaEntity, quotaRecord: QuotaRecord): Unit = { + val manager = quotaRecord.key() match { + case QuotaConfigs.CONSUMER_BYTE_RATE_OVERRIDE_CONFIG => quotaManagers.fetch + case QuotaConfigs.PRODUCER_BYTE_RATE_OVERRIDE_CONFIG => quotaManagers.produce + case QuotaConfigs.REQUEST_PERCENTAGE_OVERRIDE_CONFIG => quotaManagers.request + case QuotaConfigs.CONTROLLER_MUTATION_RATE_OVERRIDE_CONFIG => quotaManagers.controllerMutation + case _ => + warn(s"Ignoring unexpected quota key ${quotaRecord.key()} for entity $quotaEntity") + return + } + + quotaCache.updateQuotaCache(quotaEntity, quotaRecord.key, quotaRecord.value, quotaRecord.remove) + + // Convert entity into Options with sanitized values for QuotaManagers + val (sanitizedUser, sanitizedClientId) = quotaEntity match { + case UserEntity(user) => (Some(Sanitizer.sanitize(user)), None) + case DefaultUserEntity => (Some(ConfigEntityName.Default), None) + case ClientIdEntity(clientId) => (None, Some(Sanitizer.sanitize(clientId))) + case DefaultClientIdEntity => (None, Some(ConfigEntityName.Default)) + case ExplicitUserExplicitClientIdEntity(user, clientId) => (Some(Sanitizer.sanitize(user)), Some(Sanitizer.sanitize(clientId))) + case ExplicitUserDefaultClientIdEntity(user) => (Some(Sanitizer.sanitize(user)), Some(ConfigEntityName.Default)) + case DefaultUserExplicitClientIdEntity(clientId) => (Some(ConfigEntityName.Default), Some(Sanitizer.sanitize(clientId))) + case DefaultUserDefaultClientIdEntity => (Some(ConfigEntityName.Default), Some(ConfigEntityName.Default)) + case IpEntity(_) | DefaultIpEntity => throw new IllegalStateException("Should not see IP quota entities here") + } + + val quotaValue = if (quotaRecord.remove()) { + None + } else { + Some(new Quota(quotaRecord.value(), true)) + } + + manager.updateQuota( + sanitizedUser = sanitizedUser, + clientId = sanitizedClientId.map(Sanitizer.desanitize), + sanitizedClientId = sanitizedClientId, + quota = quotaValue) + } +} diff --git a/core/src/test/scala/unit/kafka/server/metadata/ClientQuotaMetadataManagerTest.scala b/core/src/test/scala/unit/kafka/server/metadata/ClientQuotaMetadataManagerTest.scala new file mode 100644 index 0000000000000..02982f44faca7 --- /dev/null +++ b/core/src/test/scala/unit/kafka/server/metadata/ClientQuotaMetadataManagerTest.scala @@ -0,0 +1,452 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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.network.ConnectionQuotas +import kafka.server.QuotaFactory.QuotaManagers +import kafka.server.{ConfigEntityName, KafkaConfig, QuotaFactory} +import kafka.utils.{MockTime, TestUtils} +import org.apache.kafka.common.config.internals.QuotaConfigs +import org.apache.kafka.common.errors.{InvalidRequestException, UnsupportedVersionException} +import org.apache.kafka.common.metadata.QuotaRecord +import org.apache.kafka.common.metrics.{Metrics, Quota} +import org.apache.kafka.common.quota.{ClientQuotaEntity, ClientQuotaFilterComponent} +import org.junit.jupiter.api.Assertions._ +import org.junit.jupiter.api.{BeforeEach, Test} +import org.mockito.ArgumentMatchers.{any, eq => _eq} +import org.mockito.Mockito._ + +import java.net.InetAddress +import java.util.Properties +import scala.collection.mutable +import scala.jdk.CollectionConverters._ + +class ClientQuotaMetadataManagerTest { + + var manager: ClientQuotaMetadataManager = _ + var cache: ClientQuotaCache = _ + + @BeforeEach + def setup(): Unit = { + val configs = TestUtils.createBrokerConfigs(1, TestUtils.MockZkConnect) + .map(KafkaConfig.fromProps(_, new Properties())) + + val time = new MockTime + val metrics = new Metrics + val quotaManagers = QuotaFactory.instantiate(configs.head, metrics, time, "quota-metadata-processor-test") + val spiedQuotaManagers = QuotaManagers( + fetch = spy(quotaManagers.fetch), + produce = spy(quotaManagers.produce), + request = spy(quotaManagers.request), + controllerMutation = spy(quotaManagers.controllerMutation), + leader = quotaManagers.leader, + follower = quotaManagers.follower, + alterLogDirs = quotaManagers.alterLogDirs, + clientQuotaCallback = quotaManagers.clientQuotaCallback + ) + val connectionQuotas = mock(classOf[ConnectionQuotas]) + cache = new ClientQuotaCache() + manager = new ClientQuotaMetadataManager(spiedQuotaManagers, connectionQuotas, cache) + } + + @Test + def testDescribeStrictMatch(): Unit = { + setupAndVerify(manager, { case (entity, _) => + val components = mutable.ListBuffer[ClientQuotaFilterComponent]() + entityToFilter(entity, components) + val results = cache.describeClientQuotas(components.toSeq, strict=true) + assertEquals(1, results.size, s"Should only match one quota for ${entity}") + }) + + val nonMatching = List( + userClientEntity("user-1", "client-id-2"), + userClientEntity("user-3", "client-id-1"), + userClientEntity("user-2", null), + userEntity("user-4"), + userClientEntity(null, "client-id-2"), + clientEntity("client-id-1"), + clientEntity("client-id-3") + ) + + nonMatching.foreach( entity => { + val components = mutable.ListBuffer[ClientQuotaFilterComponent]() + entityToFilter(entity, components) + val results = cache.describeClientQuotas(components.toSeq, strict=true) + assertEquals(0, results.size) + }) + } + + @Test + def testDescribeNonStrictMatch(): Unit = { + setupAndVerify(manager, { case (_, _) => }) + + // Match open-ended existing user. + val components = mutable.ListBuffer[ClientQuotaFilterComponent]() + entityToFilter(userEntity("user-1"), components) + var results = cache.describeClientQuotasInternal(components.toSeq, strict=false) + assertEquals(3, results.size) + assertEquals(3, results.keySet.count(quotaEntity => quotaEntity match { + case UserEntity(user) => user.equals("user-1") + case ExplicitUserDefaultClientIdEntity(user) => user.equals("user-1") + case ExplicitUserExplicitClientIdEntity(user, _) => user.equals("user-1") + case _ => false + })) + + results = cache.describeClientQuotasInternal(components.toSeq, strict=true) + assertEquals(1, results.size) + + // Match open-ended non-existent user. + components.clear() + entityToFilter(userEntity("unknown"), components) + results = cache.describeClientQuotasInternal(components.toSeq, strict=false) + assertEquals(0, results.size) + + // Match open-ended existing client ID. + components.clear() + entityToFilter(clientEntity("client-id-2"), components) + results = cache.describeClientQuotasInternal(components.toSeq, strict=false) + assertEquals(2, results.size) + assertEquals(2, results.keySet.count(quotaEntity => quotaEntity match { + case ClientIdEntity(clientId) => clientId.equals("client-id-2") + case DefaultUserExplicitClientIdEntity(clientId) => clientId.equals("client-id-2") + case ExplicitUserExplicitClientIdEntity(_, clientId) => clientId.equals("client-id-2") + case _ => false + })) + + // Match open-ended default user. + results = cache.describeClientQuotasInternal( + Seq(ClientQuotaFilterComponent.ofDefaultEntity(ClientQuotaEntity.USER)), strict=false) + assertEquals(3, results.size) + assertEquals(3, results.keySet.count(quotaEntity => quotaEntity match { + case DefaultUserEntity | DefaultUserExplicitClientIdEntity(_) | DefaultUserDefaultClientIdEntity => true + case _ => false + })) + + // Match open-ended default client. + results = cache.describeClientQuotasInternal( + Seq(ClientQuotaFilterComponent.ofDefaultEntity(ClientQuotaEntity.CLIENT_ID)), strict=false) + assertEquals(3, results.size) + assertEquals(3, results.keySet.count(quotaEntity => quotaEntity match { + case DefaultClientIdEntity | ExplicitUserDefaultClientIdEntity(_) | DefaultUserDefaultClientIdEntity => true + case _ => false + })) + } + + @Test + def testDescribeFilterOnTypes(): Unit = { + setupAndVerify(manager, { case (_, _) => }) + + var results = cache.describeClientQuotasInternal( + Seq(ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.USER)), strict=false) + assertEquals(11, results.size) + assertEquals(11, results.keySet.count(quotaEntity => quotaEntity match { + case UserEntity(_) | DefaultUserEntity | ExplicitUserExplicitClientIdEntity(_, _) | ExplicitUserDefaultClientIdEntity(_) | + DefaultUserExplicitClientIdEntity(_) | DefaultUserDefaultClientIdEntity => true + case _ => false + })) + + results = cache.describeClientQuotasInternal( + Seq(ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.CLIENT_ID)), strict=false) + assertEquals(8, results.size) + assertEquals(8, results.keySet.count(quotaEntity => quotaEntity match { + case ClientIdEntity(_) | DefaultClientIdEntity | ExplicitUserExplicitClientIdEntity(_, _) | ExplicitUserDefaultClientIdEntity(_) | + DefaultUserExplicitClientIdEntity(_) | DefaultUserDefaultClientIdEntity => true + case _ => false + })) + + results = cache.describeClientQuotasInternal( + Seq( + ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.USER), + ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.CLIENT_ID) + ), strict=true) + assertEquals(7, results.size) + assertEquals(7, results.keySet.count(quotaEntity => quotaEntity match { + case ExplicitUserExplicitClientIdEntity(_, _) | ExplicitUserDefaultClientIdEntity(_) | + DefaultUserExplicitClientIdEntity(_) | DefaultUserDefaultClientIdEntity => true + case _ => false + })) + } + + @Test + def testEntityWithDefaultName(): Unit = { + addQuotaRecord(manager, clientEntity(ConfigEntityName.Default), (QuotaConfigs.PRODUCER_BYTE_RATE_OVERRIDE_CONFIG, 20000.0)) + addQuotaRecord(manager, clientEntity(null), (QuotaConfigs.PRODUCER_BYTE_RATE_OVERRIDE_CONFIG, 30000.0)) + + val components = mutable.ListBuffer[ClientQuotaFilterComponent]() + entityToFilter(clientEntity(ConfigEntityName.Default), components) + var results = cache.describeClientQuotas(components.toSeq, strict=true) + assertEquals(1, results.size) + + components.clear() + entityToFilter(clientEntity(null), components) + results = cache.describeClientQuotas(components.toSeq, strict=true) + assertEquals(1, results.size) + } + + @Test + def testQuotaRemoval(): Unit = { + val entity = userClientEntity("user", "client-id") + addQuotaRecord(manager, entity, (QuotaConfigs.PRODUCER_BYTE_RATE_OVERRIDE_CONFIG, 10000.0)) + addQuotaRecord(manager, entity, (QuotaConfigs.CONSUMER_BYTE_RATE_OVERRIDE_CONFIG, 20000.0)) + var quotas = describeEntity(entity) + assertEquals(2, quotas.size) + assertEquals(10000.0, quotas(QuotaConfigs.PRODUCER_BYTE_RATE_OVERRIDE_CONFIG), 1e-6) + + addQuotaRecord(manager, entity, (QuotaConfigs.PRODUCER_BYTE_RATE_OVERRIDE_CONFIG, 10001.0)) + quotas = describeEntity(entity) + assertEquals(2, quotas.size) + assertEquals(10001.0, quotas(QuotaConfigs.PRODUCER_BYTE_RATE_OVERRIDE_CONFIG), 1e-6) + + addQuotaRemovalRecord(manager, entity, QuotaConfigs.PRODUCER_BYTE_RATE_OVERRIDE_CONFIG) + quotas = describeEntity(entity) + assertEquals(1, quotas.size) + assertFalse(quotas.contains(QuotaConfigs.PRODUCER_BYTE_RATE_OVERRIDE_CONFIG)) + + addQuotaRemovalRecord(manager, entity, QuotaConfigs.CONSUMER_BYTE_RATE_OVERRIDE_CONFIG) + quotas = describeEntity(entity) + assertEquals(0, quotas.size) + + // Removing non-existent quota should not do anything + addQuotaRemovalRecord(manager, entity, QuotaConfigs.CONSUMER_BYTE_RATE_OVERRIDE_CONFIG) + quotas = describeEntity(entity) + assertEquals(0, quotas.size) + } + + @Test + def testDescribeClientQuotasInvalidFilterCombination(): Unit = { + val ipFilterComponent = ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.IP) + val userFilterComponent = ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.USER) + val clientIdFilterComponent = ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.CLIENT_ID) + val expectedExceptionMessage = "Invalid entity filter component combination" + assertThrows(classOf[InvalidRequestException], + () => cache.describeClientQuotas(Seq(ipFilterComponent, userFilterComponent), strict=false), + () => expectedExceptionMessage) + assertThrows(classOf[InvalidRequestException], + () => cache.describeClientQuotas(Seq(ipFilterComponent, clientIdFilterComponent), strict=false), + () => expectedExceptionMessage) + assertThrows(classOf[InvalidRequestException], + () => cache.describeClientQuotas(Seq(ipFilterComponent, ipFilterComponent), strict = false), + () => expectedExceptionMessage) + assertThrows(classOf[InvalidRequestException], + () => cache.describeClientQuotas(Seq(userFilterComponent, userFilterComponent), strict=false), + () => expectedExceptionMessage) + } + + @Test + def testDescribeEmptyFilter(): Unit = { + var results = cache.describeClientQuotas(Seq.empty, strict=false) + assertEquals(0, results.size) + + results = cache.describeClientQuotas(Seq.empty, strict=true) + assertEquals(0, results.size) + } + + @Test + def testDescribeUnsupportedEntityType(): Unit = { + assertThrows(classOf[UnsupportedVersionException], + () => cache.describeClientQuotas(Seq(ClientQuotaFilterComponent.ofEntityType("other")), strict=false)) + } + + @Test + def testDescribeMissingEntityType(): Unit = { + assertThrows(classOf[InvalidRequestException], + () => cache.describeClientQuotas(Seq(ClientQuotaFilterComponent.ofEntity("", "name")), strict = false)) + } + + @Test + def testQuotaManagers(): Unit = { + val entity = userClientEntity("user", "client") + addQuotaRecord(manager, entity, (QuotaConfigs.CONSUMER_BYTE_RATE_OVERRIDE_CONFIG, 100.0)) + verify(manager.quotaManagers.fetch, times(1)).updateQuota( + _eq(Some("user")), + _eq(Some("client")), + _eq(Some("client")), + any(classOf[Option[Quota]]) + ) + + addQuotaRecord(manager, entity, (QuotaConfigs.PRODUCER_BYTE_RATE_OVERRIDE_CONFIG, 100.0)) + verify(manager.quotaManagers.produce, times(1)).updateQuota( + _eq(Some("user")), + _eq(Some("client")), + _eq(Some("client")), + any(classOf[Option[Quota]]) + ) + + addQuotaRecord(manager, entity, (QuotaConfigs.REQUEST_PERCENTAGE_OVERRIDE_CONFIG, 100.0)) + verify(manager.quotaManagers.request, times(1)).updateQuota( + _eq(Some("user")), + _eq(Some("client")), + _eq(Some("client")), + any(classOf[Option[Quota]]) + ) + + addQuotaRecord(manager, entity, (QuotaConfigs.CONTROLLER_MUTATION_RATE_OVERRIDE_CONFIG, 100.0)) + verify(manager.quotaManagers.controllerMutation, times(1)).updateQuota( + _eq(Some("user")), + _eq(Some("client")), + _eq(Some("client")), + any(classOf[Option[Quota]]) + ) + + addQuotaRemovalRecord(manager, entity, QuotaConfigs.CONTROLLER_MUTATION_RATE_OVERRIDE_CONFIG) + verify(manager.quotaManagers.controllerMutation, times(1)).updateQuota( + _eq(Some("user")), + _eq(Some("client")), + _eq(Some("client")), + _eq(None) + ) + } + + @Test + def testIpQuota(): Unit = { + val defaultIp = ipEntity(null) + val knownIp = ipEntity("1.2.3.4") + + addQuotaRecord(manager, defaultIp, (QuotaConfigs.IP_CONNECTION_RATE_OVERRIDE_CONFIG, 100.0)) + addQuotaRecord(manager, knownIp, (QuotaConfigs.IP_CONNECTION_RATE_OVERRIDE_CONFIG, 99.0)) + + verify(manager.connectionQuotas, times(2)).updateIpConnectionRateQuota( + any(classOf[Option[InetAddress]]), + any(classOf[Option[Int]]) + ) + + var quotas = describeEntity(defaultIp) + assertEquals(1, quotas.size) + + quotas = describeEntity(knownIp) + assertEquals(1, quotas.size) + + val results = cache.describeClientQuotas(Seq(ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.IP)), strict=false) + assertEquals(2, results.size) + + reset(manager.connectionQuotas) + addQuotaRecord(manager, knownIp, (QuotaConfigs.IP_CONNECTION_RATE_OVERRIDE_CONFIG, 98.0)) + verify(manager.connectionQuotas, times(1)).updateIpConnectionRateQuota( + any(classOf[Option[InetAddress]]), + _eq(Some(98)) + ) + + reset(manager.connectionQuotas) + addQuotaRemovalRecord(manager, knownIp, QuotaConfigs.IP_CONNECTION_RATE_OVERRIDE_CONFIG) + verify(manager.connectionQuotas, times(1)).updateIpConnectionRateQuota( + any(classOf[Option[InetAddress]]), + _eq(None) + ) + } + + @Test + def testIpQuotaUnknownKey(): Unit = { + val defaultIp = ipEntity(null) + addQuotaRecord(manager, defaultIp, ("not-an-ip-quota-key", 100.0)) + verify(manager.connectionQuotas, times(0)).updateIpConnectionRateQuota( + any(classOf[Option[InetAddress]]), + _eq(Some(100)) + ) + + assertEquals(0, describeEntity(defaultIp).size) + } + + @Test + def testUserQuotaUnknownKey(): Unit = { + val defaultUser = userEntity(null) + addQuotaRecord(manager, defaultUser, ("not-a-user-quota-key", 100.0)) + assertEquals(0, describeEntity(defaultUser).size) + } + + def setupAndVerify(manager: ClientQuotaMetadataManager, + verifier: (List[QuotaRecord.EntityData], (String, Double)) => Unit ): Unit = { + val toVerify = List( + (userClientEntity("user-1", "client-id-1"), 50.50), + (userClientEntity("user-2", "client-id-1"), 51.51), + (userClientEntity("user-3", "client-id-2"), 52.52), + (userClientEntity(null, "client-id-1"), 53.53), + (userClientEntity("user-1", null), 54.54), + (userClientEntity("user-3", null), 55.55), + (userEntity("user-1"), 56.56), + (userEntity("user-2"), 57.57), + (userEntity("user-3"), 58.58), + (userEntity(null), 59.59), + (clientEntity("client-id-2"), 60.60), + (userClientEntity(null, null), 61.61) + ) + + toVerify.foreach { + case (entity, value) => addQuotaRecord(manager, entity, (QuotaConfigs.REQUEST_PERCENTAGE_OVERRIDE_CONFIG, value)) + } + + toVerify.foreach { + case (entity, value) => verifier.apply(entity, (QuotaConfigs.REQUEST_PERCENTAGE_OVERRIDE_CONFIG, value)) + } + } + + def describeEntity(entity: List[QuotaRecord.EntityData]): Map[String, Double] = { + val components = mutable.ListBuffer[ClientQuotaFilterComponent]() + entityToFilter(entity, components) + val results = cache.describeClientQuotas(components.toSeq, strict=true) + if (results.isEmpty) { + Map() + } else if (results.size == 1) { + results.head._2 + } else { + throw new AssertionError("Matched more than one entity with strict=true describe filter") + } + } + + def addQuotaRecord(manager: ClientQuotaMetadataManager, entity: List[QuotaRecord.EntityData], quota: (String, Double)): Unit = { + manager.handleQuotaRecord(new QuotaRecord() + .setEntity(entity.asJava) + .setKey(quota._1) + .setValue(quota._2)) + } + + def addQuotaRemovalRecord(manager: ClientQuotaMetadataManager, entity: List[QuotaRecord.EntityData], quota: String): Unit = { + manager.handleQuotaRecord(new QuotaRecord() + .setEntity(entity.asJava) + .setKey(quota) + .setRemove(true)) + } + + def entityToFilter(entity: List[QuotaRecord.EntityData], components: mutable.ListBuffer[ClientQuotaFilterComponent]): Unit = { + entity.foreach(entityData => { + if (entityData.entityName() == null) { + components.append(ClientQuotaFilterComponent.ofDefaultEntity(entityData.entityType())) + } else { + components.append(ClientQuotaFilterComponent.ofEntity(entityData.entityType(), entityData.entityName())) + } + }) + } + + def clientEntity(clientId: String): List[QuotaRecord.EntityData] = { + List(new QuotaRecord.EntityData().setEntityType(ClientQuotaEntity.CLIENT_ID).setEntityName(clientId)) + } + + def userEntity(user: String): List[QuotaRecord.EntityData] = { + List(new QuotaRecord.EntityData().setEntityType(ClientQuotaEntity.USER).setEntityName(user)) + } + + def userClientEntity(user: String, clientId: String): List[QuotaRecord.EntityData] = { + List( + new QuotaRecord.EntityData().setEntityType(ClientQuotaEntity.USER).setEntityName(user), + new QuotaRecord.EntityData().setEntityType(ClientQuotaEntity.CLIENT_ID).setEntityName(clientId) + ) + } + + def ipEntity(ip: String): List[QuotaRecord.EntityData] = { + List(new QuotaRecord.EntityData().setEntityType(ClientQuotaEntity.IP).setEntityName(ip)) + } +} From 8662d0ba14bebdc22d31eefe5753fd45c67535ec Mon Sep 17 00:00:00 2001 From: David Arthur Date: Wed, 10 Feb 2021 14:26:01 -0500 Subject: [PATCH 2/2] PR feedback --- .../server/metadata/ClientQuotaCache.scala | 35 +++++++++++++------ .../metadata/ClientQuotaMetadataManager.scala | 6 ++-- 2 files changed, 27 insertions(+), 14 deletions(-) diff --git a/core/src/main/scala/kafka/server/metadata/ClientQuotaCache.scala b/core/src/main/scala/kafka/server/metadata/ClientQuotaCache.scala index aedab2e0769db..ad2378efc32cb 100644 --- a/core/src/main/scala/kafka/server/metadata/ClientQuotaCache.scala +++ b/core/src/main/scala/kafka/server/metadata/ClientQuotaCache.scala @@ -43,22 +43,35 @@ case object DefaultMatch extends QuotaMatch case object TypeMatch extends QuotaMatch +/** + * Maintains a cache of QuotaEntity and their respective quotas. + * + * The main cache is structured like: + * + * 
+ * {
+ *   (user:alice) -> {consumer_byte_rate: 10000},
+ *   (user:alice,client:x) -> {consumer_byte_rate: 8000, producer_byte_rate: 8000}
+ * }
+ *
+ * + * In addition to this cache, this class maintains three indexes for the three supported entity types (user, client, + * and IP). These indexes map a part of an entity to the list of all QuotaEntity which contain that entity. For example: + * + * 
+ * {
+ *   SpecificUser(alice) -> [(user:alice), (user:alice,client:x)],
+ *   DefaultUser -> [(user:default), (user:default, client:x)]
+ * }
+ *
+ * + * These indexes exist to support the flexible lookups needed by DescribeClientQuota RPC + */ class ClientQuotaCache { private type QuotaCacheIndex = mutable.HashMap[CacheIndexKey, mutable.HashSet[QuotaEntity]] - // A mapping of the quota entities to their quotas, for example: - // { - // (user:alice) -> {consumer_byte_rate: 10000}, - // (user:alice,client:x) -> {consumer_byte_rate: 8000, producer_byte_rate: 8000} - // } private val quotaCache = new mutable.HashMap[QuotaEntity, mutable.Map[String, Double]] - // Indexes for the three supported entity types. This is needed for flexible lookups like: "all quotas for user:alice" - // The structure of these indexes is a mapping of a specific entity entry to entities in the cache: - // { - // SpecificUser(alice) -> [(user:alice), (user:alice,client:x)], - // DefaultUser -> [(user:default), (user:default, client:default), (user:default, client:x), ...] - // } // We need three separate indexes because we also support wildcard lookups on entity type. private val userEntityIndex = new QuotaCacheIndex private val clientIdEntityIndex = new QuotaCacheIndex diff --git a/core/src/main/scala/kafka/server/metadata/ClientQuotaMetadataManager.scala b/core/src/main/scala/kafka/server/metadata/ClientQuotaMetadataManager.scala index bcbae8ae1d85a..76e0b54d90752 100644 --- a/core/src/main/scala/kafka/server/metadata/ClientQuotaMetadataManager.scala +++ b/core/src/main/scala/kafka/server/metadata/ClientQuotaMetadataManager.scala @@ -45,7 +45,7 @@ case class DefaultUserExplicitClientIdEntity(clientId: String) extends QuotaEnti case object DefaultUserDefaultClientIdEntity extends QuotaEntity /** - * Watch for changes to quotas in the metadata log and update quota managers and cache as necessary + * Process quota metadata records as they appear in the metadata log and update quota managers and cache as necessary */ class ClientQuotaMetadataManager(private[metadata] val quotaManagers: QuotaManagers, private[metadata] val connectionQuotas: ConnectionQuotas, @@ -144,8 +144,6 @@ class ClientQuotaMetadataManager(private[metadata] val quotaManagers: QuotaManag return } - quotaCache.updateQuotaCache(quotaEntity, quotaRecord.key, quotaRecord.value, quotaRecord.remove) - // Convert entity into Options with sanitized values for QuotaManagers val (sanitizedUser, sanitizedClientId) = quotaEntity match { case UserEntity(user) => (Some(Sanitizer.sanitize(user)), None) @@ -170,5 +168,7 @@ class ClientQuotaMetadataManager(private[metadata] val quotaManagers: QuotaManag clientId = sanitizedClientId.map(Sanitizer.desanitize), sanitizedClientId = sanitizedClientId, quota = quotaValue) + + quotaCache.updateQuotaCache(quotaEntity, quotaRecord.key, quotaRecord.value, quotaRecord.remove) } }