rajinisivaram commented on a change in pull request #8933:
URL: https://github.com/apache/kafka/pull/8933#discussion_r457195568
##########
File path: core/src/main/scala/kafka/server/KafkaConfig.scala
##########
@@ -864,9 +868,12 @@ object KafkaConfig {
val NumQuotaSamplesDoc = "The number of samples to retain in memory for
client quotas"
val NumReplicationQuotaSamplesDoc = "The number of samples to retain in
memory for replication quotas"
val NumAlterLogDirsReplicationQuotaSamplesDoc = "The number of samples to
retain in memory for alter log dirs replication quotas"
+ val NumControllerQuotaSamplesDoc = "The number of samples to retain in
memory for controller mutations replication quotas"
Review comment:
nit: remove `replication`
##########
File path:
core/src/test/scala/unit/kafka/server/ControllerMutationQuotaTest.scala
##########
@@ -0,0 +1,352 @@
+/**
+ * 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
+
+import java.util.Properties
+import java.util.concurrent.ExecutionException
+import java.util.concurrent.TimeUnit
+
+import kafka.utils.TestUtils
+import org.apache.kafka.common.internals.KafkaFutureImpl
+import org.apache.kafka.common.message.CreatePartitionsRequestData
+import
org.apache.kafka.common.message.CreatePartitionsRequestData.CreatePartitionsTopic
+import org.apache.kafka.common.message.CreateTopicsRequestData
+import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableTopic
+import org.apache.kafka.common.message.DeleteTopicsRequestData
+import org.apache.kafka.common.protocol.ApiKeys
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.quota.ClientQuotaAlteration
+import org.apache.kafka.common.quota.ClientQuotaEntity
+import org.apache.kafka.common.requests.AlterClientQuotasRequest
+import org.apache.kafka.common.requests.AlterClientQuotasResponse
+import org.apache.kafka.common.requests.CreatePartitionsRequest
+import org.apache.kafka.common.requests.CreatePartitionsResponse
+import org.apache.kafka.common.requests.CreateTopicsRequest
+import org.apache.kafka.common.requests.CreateTopicsResponse
+import org.apache.kafka.common.requests.DeleteTopicsRequest
+import org.apache.kafka.common.requests.DeleteTopicsResponse
+import org.apache.kafka.common.security.auth.AuthenticationContext
+import org.apache.kafka.common.security.auth.KafkaPrincipal
+import org.apache.kafka.common.security.auth.KafkaPrincipalBuilder
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertTrue
+import org.junit.Before
+import org.junit.Test
+
+import scala.jdk.CollectionConverters._
+
+object ControllerMutationQuotaTest {
+ // Principal used for all client connections. This is updated by each test.
+ var principal = KafkaPrincipal.ANONYMOUS
+ class TestPrincipalBuilder extends KafkaPrincipalBuilder {
+ override def build(context: AuthenticationContext): KafkaPrincipal = {
+ principal
+ }
+ }
+
+ def asPrincipal(newPrincipal: KafkaPrincipal)(f: => Unit): Unit = {
+ val currentPrincipal = principal
+ principal = newPrincipal
+ try f
+ finally principal = currentPrincipal
+ }
+
+ val ThrottledPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE,
"ThrottledPrincipal")
+ val UnboundedPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE,
"UnboundedPrincipal")
+
+ val StrictCreateTopicsRequestVersion = ApiKeys.CREATE_TOPICS.latestVersion
+ val PermissiveCreateTopicsRequestVersion = 5.toShort
+
+ val StrictDeleteTopicsRequestVersion = ApiKeys.DELETE_TOPICS.latestVersion
+ val PermissiveDeleteTopicsRequestVersion = 4.toShort
+
+ val StrictCreatePartitionsRequestVersion =
ApiKeys.CREATE_PARTITIONS.latestVersion
+ val PermissiveCreatePartitionsRequestVersion = 2.toShort
+
+ val TopicsWithOnePartition = Seq("topic-1" -> 1, "topic-2" -> 1)
+ val TopicsWith30Partitions = Seq("topic-1" -> 30, "topic-2" -> 30)
+ val TopicsWith31Partitions = Seq("topic-1" -> 31, "topic-2" -> 31)
+
+ val ControllerMutationRate = 2.0
+}
+
+class ControllerMutationQuotaTest extends BaseRequestTest {
+ import ControllerMutationQuotaTest._
+
+ override def brokerCount: Int = 1
+
+ override def brokerPropertyOverrides(properties: Properties): Unit = {
+ properties.put(KafkaConfig.ControlledShutdownEnableProp, "false")
+ properties.put(KafkaConfig.OffsetsTopicReplicationFactorProp, "1")
+ properties.put(KafkaConfig.OffsetsTopicPartitionsProp, "1")
+ properties.put(KafkaConfig.PrincipalBuilderClassProp,
+ classOf[ControllerMutationQuotaTest.TestPrincipalBuilder].getName)
+ // We use the default number of samples and window size.
+ properties.put(KafkaConfig.NumControllerQuotaSamplesProp, "11")
+ properties.put(KafkaConfig.ControllerQuotaWindowSizeSecondsProp, "1")
+ }
+
+ @Before
+ override def setUp(): Unit = {
+ super.setUp()
+
+ // Define a quota for ThrottledPrincipal
+ defineUserQuota(ThrottledPrincipal.getName, Some(ControllerMutationRate))
+ waitUserQuota(ThrottledPrincipal.getName, ControllerMutationRate)
+ }
+
+ @Test
+ def testSetUnsetQuota(): Unit = {
+ val principal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "User")
+ // Default Value
+ waitUserQuota(principal.getName, Long.MaxValue)
+ // Define a new quota
+ defineUserQuota(principal.getName, Some(ControllerMutationRate))
+ // Check it
+ waitUserQuota(principal.getName, ControllerMutationRate)
+ // Remove it
+ defineUserQuota(principal.getName, None)
+ // Back to the default
+ waitUserQuota(principal.getName, Long.MaxValue)
+ }
+
+ @Test
+ def testStrictCreateTopicsRequest(): Unit = {
+ asPrincipal(ThrottledPrincipal) {
+ // Create two topics worth of 30 partitions each. As we use a strict
quota, we
+ // expect the first topic to be created and the second to be rejected.
+ // Theoretically, the throttle time should be below or equal to:
+ // ((30 / 10) - 2) / 2 * 10 = 5s
+ val (throttleTimeMs1, errors1) = createTopics(TopicsWith30Partitions,
StrictCreateTopicsRequestVersion)
+ assertTrue((5000 - throttleTimeMs1) < 1000)
Review comment:
We should log the throttle time in the assertion message. Not sure if `<
1000` will result in flaky tests, will see how the PR builds in Jenkins do.
##########
File path: core/src/main/scala/kafka/server/ControllerMutationQuotaManager.scala
##########
@@ -0,0 +1,240 @@
+/**
+ * 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
+
+import kafka.network.RequestChannel
+import kafka.network.RequestChannel.Session
+import org.apache.kafka.common.MetricName
+import org.apache.kafka.common.errors.ThrottlingQuotaExceededException
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.metrics.QuotaViolationException
+import org.apache.kafka.common.metrics.Sensor
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.utils.Time
+import org.apache.kafka.server.quota.ClientQuotaCallback
+
+import scala.jdk.CollectionConverters._
+
+/**
+ * The ControllerMutationQuota trait defines a quota for a given user/clientId
pair. Such
+ * quota is not meant to be cached forever but rather during the lifetime of
processing
+ * a request.
+ */
+trait ControllerMutationQuota {
+ def isExceeded: Boolean
+ def record(permits: Double): Unit
+ def throttleTime: Int
+}
+
+/**
+ * Default quota used when quota is disabled.
+ */
+object UnboundedControllerMutationQuota extends ControllerMutationQuota {
+ override def isExceeded: Boolean = false
+ override def record(permits: Double): Unit = ()
+ override def throttleTime: Int = 0
+}
+
+/**
+ * The AbstractControllerMutationQuota is the base class of
StrictControllerMutationQuota and
+ * PermissiveControllerMutationQuota.
+ *
+ * @param time @Time object to use
+ */
+abstract class AbstractControllerMutationQuota(private val time: Time) extends
ControllerMutationQuota {
+ protected var lastThrottleTimeMs = 0L
+ protected var lastRecordedTimeMs = 0L
+
+ protected def updateThrottleTime(e: QuotaViolationException, timeMs: Long):
Unit = {
+ lastThrottleTimeMs = ClientQuotaManager.throttleTime(e, timeMs)
+ lastRecordedTimeMs = timeMs
+ }
+
+ override def throttleTime: Int = {
+ // If a throttle time has been recorded, we adjust it by deducting the
time elapsed
+ // between the recording and now. We do this because `throttleTime` may be
called
+ // long after having recorded it, especially when a request waits in the
purgatory.
+ val deltaTimeMs = time.milliseconds - lastRecordedTimeMs
+ Math.max(0, lastThrottleTimeMs - deltaTimeMs).toInt
+ }
+}
+
+/**
+ * The StrictControllerMutationQuota defines a strict quota for a given
user/clientId pair. The
+ * quota is strict meaning that it does not accept any mutations once the
quota is exhausted until
+ * it gets back to the defined rate.
Review comment:
There are two differences between Strict and Permissive quotas:
1) As described above, it does not accept any mutations if quota is already
exhausted
2) It does not throttle for any number of mutations if quota is not already
exhausted. This is different from the other one which throttles if quota will
be exceeded due to the current request.
Should we document the second behaviour as well? I am guessing 2) will be
limited by the burst value. Where would that be checked?
##########
File path: core/src/main/scala/kafka/server/ControllerMutationQuotaManager.scala
##########
@@ -0,0 +1,240 @@
+/**
+ * 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
+
+import kafka.network.RequestChannel
+import kafka.network.RequestChannel.Session
+import org.apache.kafka.common.MetricName
+import org.apache.kafka.common.errors.ThrottlingQuotaExceededException
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.metrics.QuotaViolationException
+import org.apache.kafka.common.metrics.Sensor
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.utils.Time
+import org.apache.kafka.server.quota.ClientQuotaCallback
+
+import scala.jdk.CollectionConverters._
+
+/**
+ * The ControllerMutationQuota trait defines a quota for a given user/clientId
pair. Such
+ * quota is not meant to be cached forever but rather during the lifetime of
processing
+ * a request.
+ */
+trait ControllerMutationQuota {
+ def isExceeded: Boolean
+ def record(permits: Double): Unit
+ def throttleTime: Int
+}
+
+/**
+ * Default quota used when quota is disabled.
+ */
+object UnboundedControllerMutationQuota extends ControllerMutationQuota {
+ override def isExceeded: Boolean = false
+ override def record(permits: Double): Unit = ()
+ override def throttleTime: Int = 0
+}
+
+/**
+ * The AbstractControllerMutationQuota is the base class of
StrictControllerMutationQuota and
+ * PermissiveControllerMutationQuota.
+ *
+ * @param time @Time object to use
+ */
+abstract class AbstractControllerMutationQuota(private val time: Time) extends
ControllerMutationQuota {
+ protected var lastThrottleTimeMs = 0L
+ protected var lastRecordedTimeMs = 0L
+
+ protected def updateThrottleTime(e: QuotaViolationException, timeMs: Long):
Unit = {
+ lastThrottleTimeMs = ClientQuotaManager.throttleTime(e, timeMs)
+ lastRecordedTimeMs = timeMs
+ }
+
+ override def throttleTime: Int = {
+ // If a throttle time has been recorded, we adjust it by deducting the
time elapsed
+ // between the recording and now. We do this because `throttleTime` may be
called
+ // long after having recorded it, especially when a request waits in the
purgatory.
+ val deltaTimeMs = time.milliseconds - lastRecordedTimeMs
+ Math.max(0, lastThrottleTimeMs - deltaTimeMs).toInt
+ }
+}
+
+/**
+ * The StrictControllerMutationQuota defines a strict quota for a given
user/clientId pair. The
+ * quota is strict meaning that it does not accept any mutations once the
quota is exhausted until
+ * it gets back to the defined rate.
+ *
+ * @param time @Time object to use
+ * @param quotaSensor @Sensor object with a defined quota for a given
user/clientId pair
+ */
+class StrictControllerMutationQuota(private val time: Time,
+ private val quotaSensor: Sensor)
+ extends AbstractControllerMutationQuota(time) {
+
+ override def isExceeded: Boolean = lastThrottleTimeMs > 0
+
+ override def record(permits: Double): Unit = {
+ val timeMs = time.milliseconds
+ try {
+ quotaSensor synchronized {
+ quotaSensor.checkQuotas(timeMs)
+ quotaSensor.record(permits, timeMs, false)
+ }
+ } catch {
+ case e: QuotaViolationException =>
+ updateThrottleTime(e, timeMs)
+ throw new ThrottlingQuotaExceededException(lastThrottleTimeMs.toInt,
+ Errors.THROTTLING_QUOTA_EXCEEDED.message)
+ }
+ }
+}
+
+/**
+ * The PermissiveControllerMutationQuota defines a permissive quota for a
given user/clientId pair.
+ * The quota is permissive meaning that it does accept any mutations even if
the quota is exhausted.
+ *
+ * @param time @Time object to use
+ * @param quotaSensor @Sensor object with a defined quota for a given
user/clientId pair
+ */
+class PermissiveControllerMutationQuota(private val time: Time,
+ private val quotaSensor: Sensor)
+ extends AbstractControllerMutationQuota(time) {
+
+ override def isExceeded: Boolean = false
+
+ override def record(permits: Double): Unit = {
+ val timeMs = time.milliseconds
+ try {
+ quotaSensor.record(permits, timeMs, true)
+ } catch {
+ case e: QuotaViolationException =>
+ updateThrottleTime(e, timeMs)
+ }
+ }
+}
+
+/**
+ * The ControllerMutationQuotaManager is a specialized ClientQuotaManager used
in the context
+ * of throttling controller's operations/mutations.
+ *
+ * @param config @ClientQuotaManagerConfig quota configs
+ * @param metrics @Metrics Metrics instance
+ * @param time @Time object to use
+ * @param threadNamePrefix The thread prefix to use
+ * @param quotaCallback @ClientQuotaCallback ClientQuotaCallback to use
+ */
+class ControllerMutationQuotaManager(private val config:
ClientQuotaManagerConfig,
+ private val metrics: Metrics,
+ private val time: Time,
+ private val threadNamePrefix: String,
+ private val quotaCallback:
Option[ClientQuotaCallback])
+ extends ClientQuotaManager(config, metrics, QuotaType.ControllerMutation,
time, threadNamePrefix, quotaCallback) {
+
+ override protected def clientRateMetricName(quotaMetricTags: Map[String,
String]): MetricName = {
+ metrics.metricName("mutation-rate", QuotaType.ControllerMutation.toString,
+ "Tracking mutation-rate per user/client-id",
+ quotaMetricTags.asJava)
+ }
+
+ /**
+ * Records that a user/clientId accumulated or would like to accumulate the
provided amount at the
+ * the specified time, returns throttle time in milliseconds. The quota is
strict meaning that it
+ * does not accept any mutations once the quota is exhausted until it gets
back to the defined rate.
+ *
+ * @param session The session from which the user is extracted
+ * @param clientId The client id
+ * @param value The value to accumulate
+ * @param timeMs The time at which to accumulate the value
+ * @return The throttle time in milliseconds defines as the time to wait
until the average
+ * rate gets back to the defined quota
+ */
+ override def recordAndGetThrottleTimeMs(session: Session, clientId: String,
value: Double, timeMs: Long): Int = {
+ val clientSensors = getOrCreateQuotaSensors(session, clientId)
+ val quotaSensor = clientSensors.quotaSensor
+ try {
+ quotaSensor synchronized {
+ quotaSensor.checkQuotas(timeMs)
+ quotaSensor.record(value, timeMs, false)
+ }
+ 0
+ } catch {
+ case e: QuotaViolationException =>
+ val throttleTimeMs = throttleTime(e, timeMs).toInt
+ debug(s"Quota violated for sensor (${quotaSensor.name}). Delay time:
($throttleTimeMs)")
+ throttleTimeMs
+ }
+ }
+
+ /**
+ * Returns a StrictControllerMutationQuota for the given session/clientId
pair or
+ * a UnboundedControllerMutationQuota$ if the quota is disabled.
+ *
+ * @param session The session from which the user is extracted
+ * @param clientId The client id
+ * @return ControllerMutationQuota
+ */
+ def newStrictQuotaFor(session: Session, clientId: String):
ControllerMutationQuota = {
+ if (quotasEnabled) {
+ val clientSensors = getOrCreateQuotaSensors(session, clientId)
+ new StrictControllerMutationQuota(time, clientSensors.quotaSensor)
+ } else {
+ UnboundedControllerMutationQuota
+ }
+ }
+
+ def newStrictQuotaFor(request: RequestChannel.Request):
ControllerMutationQuota =
+ newStrictQuotaFor(request.session, request.header.clientId)
+
+ /**
+ * Returns a PermissiveControllerMutationQuota for the given
session/clientId pair or
+ * a UnboundedControllerMutationQuota$ if the quota is disabled.
+ *
+ * @param session The session from which the user is extracted
+ * @param clientId The client id
+ * @return ControllerMutationQuota
+ */
+ def newPermissiveQuotaFor(session: Session, clientId: String):
ControllerMutationQuota = {
+ if (quotasEnabled) {
+ val clientSensors = getOrCreateQuotaSensors(session, clientId)
+ new PermissiveControllerMutationQuota(time, clientSensors.quotaSensor)
+ } else {
+ UnboundedControllerMutationQuota
+ }
+ }
+
+ def newPermissiveQuotaFor(request: RequestChannel.Request):
ControllerMutationQuota =
+ newPermissiveQuotaFor(request.session, request.header.clientId)
+
+ /**
+ * Returns a ControllerMutationQuota based on `strictSinceVersion`. It
returns a strict
+ * quota if the version is equal to or above of the `strictSinceVersion`, a
permissive
+ * quota if the version is bellow, and a unbounded quota if the quota is
disabled.
Review comment:
typo: below
##########
File path: core/src/main/scala/kafka/server/ControllerMutationQuotaManager.scala
##########
@@ -0,0 +1,240 @@
+/**
+ * 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
+
+import kafka.network.RequestChannel
+import kafka.network.RequestChannel.Session
+import org.apache.kafka.common.MetricName
+import org.apache.kafka.common.errors.ThrottlingQuotaExceededException
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.metrics.QuotaViolationException
+import org.apache.kafka.common.metrics.Sensor
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.utils.Time
+import org.apache.kafka.server.quota.ClientQuotaCallback
+
+import scala.jdk.CollectionConverters._
+
+/**
+ * The ControllerMutationQuota trait defines a quota for a given user/clientId
pair. Such
+ * quota is not meant to be cached forever but rather during the lifetime of
processing
+ * a request.
+ */
+trait ControllerMutationQuota {
+ def isExceeded: Boolean
+ def record(permits: Double): Unit
+ def throttleTime: Int
+}
+
+/**
+ * Default quota used when quota is disabled.
+ */
+object UnboundedControllerMutationQuota extends ControllerMutationQuota {
+ override def isExceeded: Boolean = false
+ override def record(permits: Double): Unit = ()
+ override def throttleTime: Int = 0
+}
+
+/**
+ * The AbstractControllerMutationQuota is the base class of
StrictControllerMutationQuota and
+ * PermissiveControllerMutationQuota.
+ *
+ * @param time @Time object to use
+ */
+abstract class AbstractControllerMutationQuota(private val time: Time) extends
ControllerMutationQuota {
+ protected var lastThrottleTimeMs = 0L
+ protected var lastRecordedTimeMs = 0L
+
+ protected def updateThrottleTime(e: QuotaViolationException, timeMs: Long):
Unit = {
+ lastThrottleTimeMs = ClientQuotaManager.throttleTime(e, timeMs)
+ lastRecordedTimeMs = timeMs
+ }
+
+ override def throttleTime: Int = {
+ // If a throttle time has been recorded, we adjust it by deducting the
time elapsed
+ // between the recording and now. We do this because `throttleTime` may be
called
+ // long after having recorded it, especially when a request waits in the
purgatory.
+ val deltaTimeMs = time.milliseconds - lastRecordedTimeMs
+ Math.max(0, lastThrottleTimeMs - deltaTimeMs).toInt
+ }
+}
+
+/**
+ * The StrictControllerMutationQuota defines a strict quota for a given
user/clientId pair. The
+ * quota is strict meaning that it does not accept any mutations once the
quota is exhausted until
+ * it gets back to the defined rate.
+ *
+ * @param time @Time object to use
+ * @param quotaSensor @Sensor object with a defined quota for a given
user/clientId pair
+ */
+class StrictControllerMutationQuota(private val time: Time,
+ private val quotaSensor: Sensor)
+ extends AbstractControllerMutationQuota(time) {
+
+ override def isExceeded: Boolean = lastThrottleTimeMs > 0
+
+ override def record(permits: Double): Unit = {
+ val timeMs = time.milliseconds
+ try {
+ quotaSensor synchronized {
+ quotaSensor.checkQuotas(timeMs)
+ quotaSensor.record(permits, timeMs, false)
+ }
+ } catch {
+ case e: QuotaViolationException =>
+ updateThrottleTime(e, timeMs)
+ throw new ThrottlingQuotaExceededException(lastThrottleTimeMs.toInt,
+ Errors.THROTTLING_QUOTA_EXCEEDED.message)
+ }
+ }
+}
+
+/**
+ * The PermissiveControllerMutationQuota defines a permissive quota for a
given user/clientId pair.
+ * The quota is permissive meaning that it does accept any mutations even if
the quota is exhausted.
+ *
+ * @param time @Time object to use
+ * @param quotaSensor @Sensor object with a defined quota for a given
user/clientId pair
+ */
+class PermissiveControllerMutationQuota(private val time: Time,
+ private val quotaSensor: Sensor)
+ extends AbstractControllerMutationQuota(time) {
+
+ override def isExceeded: Boolean = false
+
+ override def record(permits: Double): Unit = {
+ val timeMs = time.milliseconds
+ try {
+ quotaSensor.record(permits, timeMs, true)
+ } catch {
+ case e: QuotaViolationException =>
+ updateThrottleTime(e, timeMs)
+ }
+ }
+}
+
+/**
+ * The ControllerMutationQuotaManager is a specialized ClientQuotaManager used
in the context
+ * of throttling controller's operations/mutations.
+ *
+ * @param config @ClientQuotaManagerConfig quota configs
+ * @param metrics @Metrics Metrics instance
+ * @param time @Time object to use
+ * @param threadNamePrefix The thread prefix to use
+ * @param quotaCallback @ClientQuotaCallback ClientQuotaCallback to use
+ */
+class ControllerMutationQuotaManager(private val config:
ClientQuotaManagerConfig,
+ private val metrics: Metrics,
+ private val time: Time,
+ private val threadNamePrefix: String,
+ private val quotaCallback:
Option[ClientQuotaCallback])
+ extends ClientQuotaManager(config, metrics, QuotaType.ControllerMutation,
time, threadNamePrefix, quotaCallback) {
+
+ override protected def clientRateMetricName(quotaMetricTags: Map[String,
String]): MetricName = {
+ metrics.metricName("mutation-rate", QuotaType.ControllerMutation.toString,
+ "Tracking mutation-rate per user/client-id",
+ quotaMetricTags.asJava)
+ }
+
+ /**
+ * Records that a user/clientId accumulated or would like to accumulate the
provided amount at the
+ * the specified time, returns throttle time in milliseconds. The quota is
strict meaning that it
+ * does not accept any mutations once the quota is exhausted until it gets
back to the defined rate.
+ *
+ * @param session The session from which the user is extracted
+ * @param clientId The client id
+ * @param value The value to accumulate
+ * @param timeMs The time at which to accumulate the value
+ * @return The throttle time in milliseconds defines as the time to wait
until the average
+ * rate gets back to the defined quota
+ */
+ override def recordAndGetThrottleTimeMs(session: Session, clientId: String,
value: Double, timeMs: Long): Int = {
+ val clientSensors = getOrCreateQuotaSensors(session, clientId)
+ val quotaSensor = clientSensors.quotaSensor
+ try {
+ quotaSensor synchronized {
+ quotaSensor.checkQuotas(timeMs)
+ quotaSensor.record(value, timeMs, false)
+ }
+ 0
+ } catch {
+ case e: QuotaViolationException =>
+ val throttleTimeMs = throttleTime(e, timeMs).toInt
+ debug(s"Quota violated for sensor (${quotaSensor.name}). Delay time:
($throttleTimeMs)")
+ throttleTimeMs
+ }
+ }
+
+ /**
+ * Returns a StrictControllerMutationQuota for the given session/clientId
pair or
+ * a UnboundedControllerMutationQuota$ if the quota is disabled.
+ *
+ * @param session The session from which the user is extracted
+ * @param clientId The client id
+ * @return ControllerMutationQuota
+ */
+ def newStrictQuotaFor(session: Session, clientId: String):
ControllerMutationQuota = {
+ if (quotasEnabled) {
+ val clientSensors = getOrCreateQuotaSensors(session, clientId)
+ new StrictControllerMutationQuota(time, clientSensors.quotaSensor)
+ } else {
+ UnboundedControllerMutationQuota
+ }
+ }
+
+ def newStrictQuotaFor(request: RequestChannel.Request):
ControllerMutationQuota =
+ newStrictQuotaFor(request.session, request.header.clientId)
+
+ /**
+ * Returns a PermissiveControllerMutationQuota for the given
session/clientId pair or
Review comment:
nit: user/clientId pair?
##########
File path: core/src/main/scala/kafka/server/ReplicationQuotaManager.scala
##########
@@ -20,13 +20,11 @@ import java.util.concurrent.{ConcurrentHashMap, TimeUnit}
import java.util.concurrent.locks.ReentrantReadWriteLock
import scala.collection.Seq
-
Review comment:
Revert changes to file since nothing has actually changed?
##########
File path: core/src/main/scala/kafka/server/ControllerMutationQuotaManager.scala
##########
@@ -0,0 +1,240 @@
+/**
+ * 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
+
+import kafka.network.RequestChannel
+import kafka.network.RequestChannel.Session
+import org.apache.kafka.common.MetricName
+import org.apache.kafka.common.errors.ThrottlingQuotaExceededException
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.metrics.QuotaViolationException
+import org.apache.kafka.common.metrics.Sensor
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.utils.Time
+import org.apache.kafka.server.quota.ClientQuotaCallback
+
+import scala.jdk.CollectionConverters._
+
+/**
+ * The ControllerMutationQuota trait defines a quota for a given user/clientId
pair. Such
+ * quota is not meant to be cached forever but rather during the lifetime of
processing
+ * a request.
+ */
+trait ControllerMutationQuota {
+ def isExceeded: Boolean
+ def record(permits: Double): Unit
+ def throttleTime: Int
+}
+
+/**
+ * Default quota used when quota is disabled.
+ */
+object UnboundedControllerMutationQuota extends ControllerMutationQuota {
+ override def isExceeded: Boolean = false
+ override def record(permits: Double): Unit = ()
+ override def throttleTime: Int = 0
+}
+
+/**
+ * The AbstractControllerMutationQuota is the base class of
StrictControllerMutationQuota and
+ * PermissiveControllerMutationQuota.
+ *
+ * @param time @Time object to use
+ */
+abstract class AbstractControllerMutationQuota(private val time: Time) extends
ControllerMutationQuota {
+ protected var lastThrottleTimeMs = 0L
+ protected var lastRecordedTimeMs = 0L
+
+ protected def updateThrottleTime(e: QuotaViolationException, timeMs: Long):
Unit = {
+ lastThrottleTimeMs = ClientQuotaManager.throttleTime(e, timeMs)
+ lastRecordedTimeMs = timeMs
+ }
+
+ override def throttleTime: Int = {
+ // If a throttle time has been recorded, we adjust it by deducting the
time elapsed
+ // between the recording and now. We do this because `throttleTime` may be
called
+ // long after having recorded it, especially when a request waits in the
purgatory.
+ val deltaTimeMs = time.milliseconds - lastRecordedTimeMs
+ Math.max(0, lastThrottleTimeMs - deltaTimeMs).toInt
+ }
+}
+
+/**
+ * The StrictControllerMutationQuota defines a strict quota for a given
user/clientId pair. The
+ * quota is strict meaning that it does not accept any mutations once the
quota is exhausted until
+ * it gets back to the defined rate.
+ *
+ * @param time @Time object to use
+ * @param quotaSensor @Sensor object with a defined quota for a given
user/clientId pair
+ */
+class StrictControllerMutationQuota(private val time: Time,
+ private val quotaSensor: Sensor)
+ extends AbstractControllerMutationQuota(time) {
+
+ override def isExceeded: Boolean = lastThrottleTimeMs > 0
+
+ override def record(permits: Double): Unit = {
+ val timeMs = time.milliseconds
+ try {
+ quotaSensor synchronized {
+ quotaSensor.checkQuotas(timeMs)
+ quotaSensor.record(permits, timeMs, false)
+ }
+ } catch {
+ case e: QuotaViolationException =>
+ updateThrottleTime(e, timeMs)
+ throw new ThrottlingQuotaExceededException(lastThrottleTimeMs.toInt,
+ Errors.THROTTLING_QUOTA_EXCEEDED.message)
+ }
+ }
+}
+
+/**
+ * The PermissiveControllerMutationQuota defines a permissive quota for a
given user/clientId pair.
+ * The quota is permissive meaning that it does accept any mutations even if
the quota is exhausted.
+ *
+ * @param time @Time object to use
+ * @param quotaSensor @Sensor object with a defined quota for a given
user/clientId pair
+ */
+class PermissiveControllerMutationQuota(private val time: Time,
+ private val quotaSensor: Sensor)
+ extends AbstractControllerMutationQuota(time) {
+
+ override def isExceeded: Boolean = false
+
+ override def record(permits: Double): Unit = {
+ val timeMs = time.milliseconds
+ try {
+ quotaSensor.record(permits, timeMs, true)
+ } catch {
+ case e: QuotaViolationException =>
+ updateThrottleTime(e, timeMs)
+ }
+ }
+}
+
+/**
+ * The ControllerMutationQuotaManager is a specialized ClientQuotaManager used
in the context
+ * of throttling controller's operations/mutations.
+ *
+ * @param config @ClientQuotaManagerConfig quota configs
+ * @param metrics @Metrics Metrics instance
+ * @param time @Time object to use
+ * @param threadNamePrefix The thread prefix to use
+ * @param quotaCallback @ClientQuotaCallback ClientQuotaCallback to use
+ */
+class ControllerMutationQuotaManager(private val config:
ClientQuotaManagerConfig,
+ private val metrics: Metrics,
+ private val time: Time,
+ private val threadNamePrefix: String,
+ private val quotaCallback:
Option[ClientQuotaCallback])
+ extends ClientQuotaManager(config, metrics, QuotaType.ControllerMutation,
time, threadNamePrefix, quotaCallback) {
+
+ override protected def clientRateMetricName(quotaMetricTags: Map[String,
String]): MetricName = {
+ metrics.metricName("mutation-rate", QuotaType.ControllerMutation.toString,
+ "Tracking mutation-rate per user/client-id",
+ quotaMetricTags.asJava)
+ }
+
+ /**
+ * Records that a user/clientId accumulated or would like to accumulate the
provided amount at the
+ * the specified time, returns throttle time in milliseconds. The quota is
strict meaning that it
+ * does not accept any mutations once the quota is exhausted until it gets
back to the defined rate.
+ *
+ * @param session The session from which the user is extracted
+ * @param clientId The client id
+ * @param value The value to accumulate
+ * @param timeMs The time at which to accumulate the value
+ * @return The throttle time in milliseconds defines as the time to wait
until the average
+ * rate gets back to the defined quota
+ */
+ override def recordAndGetThrottleTimeMs(session: Session, clientId: String,
value: Double, timeMs: Long): Int = {
+ val clientSensors = getOrCreateQuotaSensors(session, clientId)
+ val quotaSensor = clientSensors.quotaSensor
+ try {
+ quotaSensor synchronized {
+ quotaSensor.checkQuotas(timeMs)
+ quotaSensor.record(value, timeMs, false)
+ }
+ 0
+ } catch {
+ case e: QuotaViolationException =>
+ val throttleTimeMs = throttleTime(e, timeMs).toInt
+ debug(s"Quota violated for sensor (${quotaSensor.name}). Delay time:
($throttleTimeMs)")
+ throttleTimeMs
+ }
+ }
+
+ /**
+ * Returns a StrictControllerMutationQuota for the given session/clientId
pair or
Review comment:
nit: user/clientId pair?
##########
File path: core/src/main/scala/kafka/server/ControllerMutationQuotaManager.scala
##########
@@ -0,0 +1,240 @@
+/**
+ * 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
+
+import kafka.network.RequestChannel
+import kafka.network.RequestChannel.Session
+import org.apache.kafka.common.MetricName
+import org.apache.kafka.common.errors.ThrottlingQuotaExceededException
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.metrics.QuotaViolationException
+import org.apache.kafka.common.metrics.Sensor
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.utils.Time
+import org.apache.kafka.server.quota.ClientQuotaCallback
+
+import scala.jdk.CollectionConverters._
+
+/**
+ * The ControllerMutationQuota trait defines a quota for a given user/clientId
pair. Such
+ * quota is not meant to be cached forever but rather during the lifetime of
processing
+ * a request.
+ */
+trait ControllerMutationQuota {
+ def isExceeded: Boolean
+ def record(permits: Double): Unit
+ def throttleTime: Int
+}
+
+/**
+ * Default quota used when quota is disabled.
+ */
+object UnboundedControllerMutationQuota extends ControllerMutationQuota {
+ override def isExceeded: Boolean = false
+ override def record(permits: Double): Unit = ()
+ override def throttleTime: Int = 0
+}
+
+/**
+ * The AbstractControllerMutationQuota is the base class of
StrictControllerMutationQuota and
+ * PermissiveControllerMutationQuota.
+ *
+ * @param time @Time object to use
+ */
+abstract class AbstractControllerMutationQuota(private val time: Time) extends
ControllerMutationQuota {
+ protected var lastThrottleTimeMs = 0L
+ protected var lastRecordedTimeMs = 0L
+
+ protected def updateThrottleTime(e: QuotaViolationException, timeMs: Long):
Unit = {
+ lastThrottleTimeMs = ClientQuotaManager.throttleTime(e, timeMs)
+ lastRecordedTimeMs = timeMs
+ }
+
+ override def throttleTime: Int = {
+ // If a throttle time has been recorded, we adjust it by deducting the
time elapsed
+ // between the recording and now. We do this because `throttleTime` may be
called
+ // long after having recorded it, especially when a request waits in the
purgatory.
+ val deltaTimeMs = time.milliseconds - lastRecordedTimeMs
+ Math.max(0, lastThrottleTimeMs - deltaTimeMs).toInt
+ }
+}
+
+/**
+ * The StrictControllerMutationQuota defines a strict quota for a given
user/clientId pair. The
+ * quota is strict meaning that it does not accept any mutations once the
quota is exhausted until
+ * it gets back to the defined rate.
+ *
+ * @param time @Time object to use
+ * @param quotaSensor @Sensor object with a defined quota for a given
user/clientId pair
+ */
+class StrictControllerMutationQuota(private val time: Time,
+ private val quotaSensor: Sensor)
+ extends AbstractControllerMutationQuota(time) {
+
+ override def isExceeded: Boolean = lastThrottleTimeMs > 0
+
+ override def record(permits: Double): Unit = {
+ val timeMs = time.milliseconds
+ try {
+ quotaSensor synchronized {
+ quotaSensor.checkQuotas(timeMs)
+ quotaSensor.record(permits, timeMs, false)
+ }
+ } catch {
+ case e: QuotaViolationException =>
+ updateThrottleTime(e, timeMs)
+ throw new ThrottlingQuotaExceededException(lastThrottleTimeMs.toInt,
+ Errors.THROTTLING_QUOTA_EXCEEDED.message)
+ }
+ }
+}
+
+/**
+ * The PermissiveControllerMutationQuota defines a permissive quota for a
given user/clientId pair.
+ * The quota is permissive meaning that it does accept any mutations even if
the quota is exhausted.
+ *
+ * @param time @Time object to use
+ * @param quotaSensor @Sensor object with a defined quota for a given
user/clientId pair
+ */
+class PermissiveControllerMutationQuota(private val time: Time,
+ private val quotaSensor: Sensor)
+ extends AbstractControllerMutationQuota(time) {
+
+ override def isExceeded: Boolean = false
+
+ override def record(permits: Double): Unit = {
+ val timeMs = time.milliseconds
+ try {
+ quotaSensor.record(permits, timeMs, true)
+ } catch {
+ case e: QuotaViolationException =>
+ updateThrottleTime(e, timeMs)
+ }
+ }
+}
+
+/**
+ * The ControllerMutationQuotaManager is a specialized ClientQuotaManager used
in the context
+ * of throttling controller's operations/mutations.
+ *
+ * @param config @ClientQuotaManagerConfig quota configs
+ * @param metrics @Metrics Metrics instance
+ * @param time @Time object to use
+ * @param threadNamePrefix The thread prefix to use
+ * @param quotaCallback @ClientQuotaCallback ClientQuotaCallback to use
+ */
+class ControllerMutationQuotaManager(private val config:
ClientQuotaManagerConfig,
+ private val metrics: Metrics,
+ private val time: Time,
+ private val threadNamePrefix: String,
+ private val quotaCallback:
Option[ClientQuotaCallback])
+ extends ClientQuotaManager(config, metrics, QuotaType.ControllerMutation,
time, threadNamePrefix, quotaCallback) {
+
+ override protected def clientRateMetricName(quotaMetricTags: Map[String,
String]): MetricName = {
+ metrics.metricName("mutation-rate", QuotaType.ControllerMutation.toString,
+ "Tracking mutation-rate per user/client-id",
+ quotaMetricTags.asJava)
+ }
+
+ /**
+ * Records that a user/clientId accumulated or would like to accumulate the
provided amount at the
+ * the specified time, returns throttle time in milliseconds. The quota is
strict meaning that it
+ * does not accept any mutations once the quota is exhausted until it gets
back to the defined rate.
+ *
+ * @param session The session from which the user is extracted
+ * @param clientId The client id
+ * @param value The value to accumulate
+ * @param timeMs The time at which to accumulate the value
+ * @return The throttle time in milliseconds defines as the time to wait
until the average
+ * rate gets back to the defined quota
+ */
+ override def recordAndGetThrottleTimeMs(session: Session, clientId: String,
value: Double, timeMs: Long): Int = {
+ val clientSensors = getOrCreateQuotaSensors(session, clientId)
+ val quotaSensor = clientSensors.quotaSensor
+ try {
+ quotaSensor synchronized {
+ quotaSensor.checkQuotas(timeMs)
+ quotaSensor.record(value, timeMs, false)
+ }
+ 0
+ } catch {
+ case e: QuotaViolationException =>
+ val throttleTimeMs = throttleTime(e, timeMs).toInt
+ debug(s"Quota violated for sensor (${quotaSensor.name}). Delay time:
($throttleTimeMs)")
+ throttleTimeMs
+ }
+ }
+
+ /**
+ * Returns a StrictControllerMutationQuota for the given session/clientId
pair or
+ * a UnboundedControllerMutationQuota$ if the quota is disabled.
+ *
+ * @param session The session from which the user is extracted
+ * @param clientId The client id
+ * @return ControllerMutationQuota
+ */
+ def newStrictQuotaFor(session: Session, clientId: String):
ControllerMutationQuota = {
+ if (quotasEnabled) {
+ val clientSensors = getOrCreateQuotaSensors(session, clientId)
+ new StrictControllerMutationQuota(time, clientSensors.quotaSensor)
+ } else {
+ UnboundedControllerMutationQuota
+ }
+ }
+
+ def newStrictQuotaFor(request: RequestChannel.Request):
ControllerMutationQuota =
+ newStrictQuotaFor(request.session, request.header.clientId)
+
+ /**
+ * Returns a PermissiveControllerMutationQuota for the given
session/clientId pair or
+ * a UnboundedControllerMutationQuota$ if the quota is disabled.
+ *
+ * @param session The session from which the user is extracted
+ * @param clientId The client id
+ * @return ControllerMutationQuota
+ */
+ def newPermissiveQuotaFor(session: Session, clientId: String):
ControllerMutationQuota = {
+ if (quotasEnabled) {
+ val clientSensors = getOrCreateQuotaSensors(session, clientId)
+ new PermissiveControllerMutationQuota(time, clientSensors.quotaSensor)
+ } else {
+ UnboundedControllerMutationQuota
+ }
+ }
+
+ def newPermissiveQuotaFor(request: RequestChannel.Request):
ControllerMutationQuota =
+ newPermissiveQuotaFor(request.session, request.header.clientId)
+
+ /**
+ * Returns a ControllerMutationQuota based on `strictSinceVersion`. It
returns a strict
+ * quota if the version is equal to or above of the `strictSinceVersion`, a
permissive
+ * quota if the version is bellow, and a unbounded quota if the quota is
disabled.
+ *
+ * When the quota is strictly enforced. Any operation above the quota is not
allowed
+ * and rejected with a THROTTLING_QUOTA_EXCEEDED error.
+ *
+ * @param request The request to extract the session and the clientId from
Review comment:
nit: user and clientId of the session
##########
File path:
core/src/test/scala/unit/kafka/server/ControllerMutationQuotaTest.scala
##########
@@ -0,0 +1,352 @@
+/**
+ * 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
+
+import java.util.Properties
+import java.util.concurrent.ExecutionException
+import java.util.concurrent.TimeUnit
+
+import kafka.utils.TestUtils
+import org.apache.kafka.common.internals.KafkaFutureImpl
+import org.apache.kafka.common.message.CreatePartitionsRequestData
+import
org.apache.kafka.common.message.CreatePartitionsRequestData.CreatePartitionsTopic
+import org.apache.kafka.common.message.CreateTopicsRequestData
+import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableTopic
+import org.apache.kafka.common.message.DeleteTopicsRequestData
+import org.apache.kafka.common.protocol.ApiKeys
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.quota.ClientQuotaAlteration
+import org.apache.kafka.common.quota.ClientQuotaEntity
+import org.apache.kafka.common.requests.AlterClientQuotasRequest
+import org.apache.kafka.common.requests.AlterClientQuotasResponse
+import org.apache.kafka.common.requests.CreatePartitionsRequest
+import org.apache.kafka.common.requests.CreatePartitionsResponse
+import org.apache.kafka.common.requests.CreateTopicsRequest
+import org.apache.kafka.common.requests.CreateTopicsResponse
+import org.apache.kafka.common.requests.DeleteTopicsRequest
+import org.apache.kafka.common.requests.DeleteTopicsResponse
+import org.apache.kafka.common.security.auth.AuthenticationContext
+import org.apache.kafka.common.security.auth.KafkaPrincipal
+import org.apache.kafka.common.security.auth.KafkaPrincipalBuilder
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertTrue
+import org.junit.Before
+import org.junit.Test
+
+import scala.jdk.CollectionConverters._
+
+object ControllerMutationQuotaTest {
+ // Principal used for all client connections. This is updated by each test.
+ var principal = KafkaPrincipal.ANONYMOUS
+ class TestPrincipalBuilder extends KafkaPrincipalBuilder {
+ override def build(context: AuthenticationContext): KafkaPrincipal = {
+ principal
+ }
+ }
+
+ def asPrincipal(newPrincipal: KafkaPrincipal)(f: => Unit): Unit = {
+ val currentPrincipal = principal
+ principal = newPrincipal
+ try f
+ finally principal = currentPrincipal
+ }
+
+ val ThrottledPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE,
"ThrottledPrincipal")
+ val UnboundedPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE,
"UnboundedPrincipal")
+
+ val StrictCreateTopicsRequestVersion = ApiKeys.CREATE_TOPICS.latestVersion
+ val PermissiveCreateTopicsRequestVersion = 5.toShort
+
+ val StrictDeleteTopicsRequestVersion = ApiKeys.DELETE_TOPICS.latestVersion
+ val PermissiveDeleteTopicsRequestVersion = 4.toShort
+
+ val StrictCreatePartitionsRequestVersion =
ApiKeys.CREATE_PARTITIONS.latestVersion
+ val PermissiveCreatePartitionsRequestVersion = 2.toShort
+
+ val TopicsWithOnePartition = Seq("topic-1" -> 1, "topic-2" -> 1)
+ val TopicsWith30Partitions = Seq("topic-1" -> 30, "topic-2" -> 30)
+ val TopicsWith31Partitions = Seq("topic-1" -> 31, "topic-2" -> 31)
+
+ val ControllerMutationRate = 2.0
+}
+
+class ControllerMutationQuotaTest extends BaseRequestTest {
+ import ControllerMutationQuotaTest._
+
+ override def brokerCount: Int = 1
+
+ override def brokerPropertyOverrides(properties: Properties): Unit = {
+ properties.put(KafkaConfig.ControlledShutdownEnableProp, "false")
+ properties.put(KafkaConfig.OffsetsTopicReplicationFactorProp, "1")
+ properties.put(KafkaConfig.OffsetsTopicPartitionsProp, "1")
+ properties.put(KafkaConfig.PrincipalBuilderClassProp,
+ classOf[ControllerMutationQuotaTest.TestPrincipalBuilder].getName)
+ // We use the default number of samples and window size.
+ properties.put(KafkaConfig.NumControllerQuotaSamplesProp, "11")
+ properties.put(KafkaConfig.ControllerQuotaWindowSizeSecondsProp, "1")
+ }
+
+ @Before
+ override def setUp(): Unit = {
+ super.setUp()
+
+ // Define a quota for ThrottledPrincipal
+ defineUserQuota(ThrottledPrincipal.getName, Some(ControllerMutationRate))
+ waitUserQuota(ThrottledPrincipal.getName, ControllerMutationRate)
+ }
+
+ @Test
+ def testSetUnsetQuota(): Unit = {
+ val principal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "User")
+ // Default Value
+ waitUserQuota(principal.getName, Long.MaxValue)
+ // Define a new quota
+ defineUserQuota(principal.getName, Some(ControllerMutationRate))
+ // Check it
+ waitUserQuota(principal.getName, ControllerMutationRate)
+ // Remove it
+ defineUserQuota(principal.getName, None)
+ // Back to the default
+ waitUserQuota(principal.getName, Long.MaxValue)
+ }
+
+ @Test
+ def testStrictCreateTopicsRequest(): Unit = {
+ asPrincipal(ThrottledPrincipal) {
+ // Create two topics worth of 30 partitions each. As we use a strict
quota, we
+ // expect the first topic to be created and the second to be rejected.
+ // Theoretically, the throttle time should be below or equal to:
+ // ((30 / 10) - 2) / 2 * 10 = 5s
+ val (throttleTimeMs1, errors1) = createTopics(TopicsWith30Partitions,
StrictCreateTopicsRequestVersion)
+ assertTrue((5000 - throttleTimeMs1) < 1000)
+ assertEquals(Seq(Errors.NONE, Errors.THROTTLING_QUOTA_EXCEEDED), errors1)
+
+ // Retry the second topic. It should succeed after the throttling delay
is passed and the
+ // throttle time should be zero.
+ TestUtils.waitUntilTrue(() => {
+ val (throttleTimeMs2, errors2) =
createTopics(TopicsWith30Partitions.drop(1), StrictCreateTopicsRequestVersion)
+ throttleTimeMs2 == 0 && errors2 == Seq(Errors.NONE)
+ }, "Failed to create topics after having been throttled")
+ }
+ }
+
+ @Test
+ def testPermissiveCreateTopicsRequest(): Unit = {
+ asPrincipal(ThrottledPrincipal) {
+ // Create two topics worth of 30 partitions each. As we use a permissive
quota, we
+ // expect both topics to be created.
+ // Theoretically, the throttle time should be below or equal to:
+ // ((60 / 10) - 2) / 2 * 10 = 20s
+ val (throttleTimeMs, errors) = createTopics(TopicsWith30Partitions,
PermissiveCreateTopicsRequestVersion)
+ assertTrue((20000 - throttleTimeMs) < 1000)
Review comment:
nit: Include actual `throttleTimeMs` in assertion message (multiple
places)
##########
File path:
core/src/test/scala/unit/kafka/server/ThrottledChannelExpirationTest.scala
##########
@@ -14,10 +14,8 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-
Review comment:
Seems to be changing a file for no reason
##########
File path: core/src/main/scala/kafka/server/MetadataCache.scala
##########
@@ -222,6 +222,10 @@ class MetadataCache(brokerId: Int) extends Logging {
metadataSnapshot.partitionStates.get(topic).flatMap(_.get(partitionId))
}
+ def numPartitions(topic: String): Int = {
+ metadataSnapshot.partitionStates.get(topic).map(_.size).getOrElse(0)
Review comment:
nit: return Option[Int] and the caller can do `getOrElse(0)`?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
