This is an automated email from the ASF dual-hosted git repository.
chia7712 pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git
The following commit(s) were added to refs/heads/trunk by this push:
new 4fc9c58ba4a KAFKA-17840 Move ClientRequestQuotaManager to server
module (#22403)
4fc9c58ba4a is described below
commit 4fc9c58ba4a40e647f9acc0cb9744749e4d2dff8
Author: Mickael Maison <[email protected]>
AuthorDate: Wed Jun 24 22:27:45 2026 +0200
KAFKA-17840 Move ClientRequestQuotaManager to server module (#22403)
Also move the quota manager tests classes
Reviewers: Chia-Ping Tsai <[email protected]>
---
core/src/main/java/kafka/server/QuotaFactory.java | 1 +
.../kafka/server/BaseClientQuotaManagerTest.scala | 84 ---
.../unit/kafka/server/ClientQuotaManagerTest.scala | 626 ---------------------
.../server/ClientRequestQuotaManagerTest.scala | 99 ----
.../unit/kafka/server/ControllerApisTest.scala | 2 +-
.../ControllerMutationQuotaManagerTest.scala | 243 --------
.../scala/unit/kafka/server/KafkaApisTest.scala | 2 +-
.../metadata/KRaftMetadataRequestBenchmark.java | 2 +-
.../server/quota}/ClientRequestQuotaManager.java | 6 +-
.../server/quota/BaseClientQuotaManagerTest.java | 97 ++++
.../kafka/server/quota/ClientQuotaManagerTest.java | 620 ++++++++++++++++++++
.../quota/ClientRequestQuotaManagerTest.java | 104 ++++
.../quota/ControllerMutationQuotaManagerTest.java | 243 ++++++++
13 files changed, 1069 insertions(+), 1060 deletions(-)
diff --git a/core/src/main/java/kafka/server/QuotaFactory.java
b/core/src/main/java/kafka/server/QuotaFactory.java
index aef041dff46..aa28856cdbb 100644
--- a/core/src/main/java/kafka/server/QuotaFactory.java
+++ b/core/src/main/java/kafka/server/QuotaFactory.java
@@ -27,6 +27,7 @@ import org.apache.kafka.server.config.QuotaConfig;
import org.apache.kafka.server.config.ReplicationQuotaManagerConfig;
import org.apache.kafka.server.quota.ClientQuotaCallback;
import org.apache.kafka.server.quota.ClientQuotaManager;
+import org.apache.kafka.server.quota.ClientRequestQuotaManager;
import org.apache.kafka.server.quota.ControllerMutationQuotaManager;
import org.apache.kafka.server.quota.QuotaType;
import org.apache.kafka.server.quota.ReplicaQuota;
diff --git
a/core/src/test/scala/unit/kafka/server/BaseClientQuotaManagerTest.scala
b/core/src/test/scala/unit/kafka/server/BaseClientQuotaManagerTest.scala
deleted file mode 100644
index 294d5f27988..00000000000
--- a/core/src/test/scala/unit/kafka/server/BaseClientQuotaManagerTest.scala
+++ /dev/null
@@ -1,84 +0,0 @@
-/**
- * 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 java.net.InetAddress
-import java.util
-import java.util.Collections
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.memory.MemoryPool
-import org.apache.kafka.common.metrics.{MetricConfig, Metrics}
-import org.apache.kafka.common.network.{ClientInformation, ListenerName}
-import org.apache.kafka.common.protocol.ApiKeys
-import org.apache.kafka.common.requests.FetchRequest.PartitionData
-import org.apache.kafka.common.requests.{AbstractRequest, FetchRequest,
RequestContext, RequestHeader}
-import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
-import org.apache.kafka.common.utils.MockTime
-import org.apache.kafka.network.{Request, Session}
-import org.apache.kafka.network.metrics.RequestChannelMetrics
-import org.apache.kafka.server.quota.{ClientQuotaManager, ThrottleCallback}
-import org.junit.jupiter.api.AfterEach
-import org.mockito.Mockito.mock
-
-class BaseClientQuotaManagerTest {
- protected val time = new MockTime
- protected var numCallbacks: Int = 0
- protected val metrics = new Metrics(new MetricConfig(),
Collections.emptyList(), time)
-
- @AfterEach
- def tearDown(): Unit = {
- metrics.close()
- }
-
- protected def callback: ThrottleCallback = new ThrottleCallback {
- override def startThrottling(): Unit = {}
- override def endThrottling(): Unit = {
- // Count how many times this callback is called for
notifyThrottlingDone().
- numCallbacks += 1
- }
- }
-
- protected def buildRequest[T <: AbstractRequest](builder:
AbstractRequest.Builder[T],
- listenerName: ListenerName
= ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)): (T, Request) =
{
-
- val request = builder.build()
- val buffer = request.serializeWithHeader(new RequestHeader(builder.apiKey,
request.version, "", 0))
- val requestChannelMetrics: RequestChannelMetrics =
mock(classOf[RequestChannelMetrics])
-
- // read the header from the buffer first so that the body can be read next
from the Request constructor
- val header = RequestHeader.parse(buffer)
- val context = new RequestContext(header, "1", InetAddress.getLocalHost,
KafkaPrincipal.ANONYMOUS,
- listenerName, SecurityProtocol.PLAINTEXT, ClientInformation.EMPTY, false)
- (request, new Request(1, context, 0, MemoryPool.NONE, buffer,
requestChannelMetrics))
- }
-
- protected def buildSession(user: String): Session = {
- val principal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, user)
- new Session(principal, null)
- }
-
- protected def maybeRecord(quotaManager: ClientQuotaManager, user: String,
clientId: String, value: Double): Int = {
-
- quotaManager.maybeRecordAndGetThrottleTimeMs(buildSession(user), clientId,
value, time.milliseconds)
- }
-
- protected def throttle(quotaManager: ClientQuotaManager, user: String,
clientId: String, throttleTimeMs: Int,
- channelThrottlingCallback: ThrottleCallback): Unit = {
- val (_, request) =
buildRequest(FetchRequest.Builder.forConsumer(ApiKeys.FETCH.latestVersion, 0,
1000, new util.HashMap[TopicPartition, PartitionData]))
- quotaManager.throttle(request.header.clientId(), request.session,
channelThrottlingCallback, throttleTimeMs)
- }
-}
diff --git a/core/src/test/scala/unit/kafka/server/ClientQuotaManagerTest.scala
b/core/src/test/scala/unit/kafka/server/ClientQuotaManagerTest.scala
deleted file mode 100644
index f0d86ba9fdf..00000000000
--- a/core/src/test/scala/unit/kafka/server/ClientQuotaManagerTest.scala
+++ /dev/null
@@ -1,626 +0,0 @@
-/**
- * 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 java.net.InetAddress
-import org.apache.kafka.common.internals.Plugin
-import org.apache.kafka.common.metrics.Quota
-import org.apache.kafka.common.security.auth.KafkaPrincipal
-import org.apache.kafka.network.Session
-import org.apache.kafka.server.config.ClientQuotaManagerConfig
-import org.apache.kafka.server.quota.{ClientQuotaCallback, ClientQuotaEntity,
ClientQuotaManager, ClientQuotaType, QuotaType}
-import org.junit.jupiter.api.Assertions._
-import org.junit.jupiter.api.Test
-
-import java.util.{Collections, Map, HashMap, Optional}
-
-class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
- private val config = new ClientQuotaManagerConfig()
-
- private def testQuotaParsing(config: ClientQuotaManagerConfig, client1:
UserClient, client2: UserClient, randomClient: UserClient, defaultConfigClient:
UserClient): Unit = {
- val clientQuotaManager = new ClientQuotaManager(config, metrics,
QuotaType.PRODUCE, time, "")
-
- try {
- // Case 1: Update the quota. Assert that the new quota value is returned
- clientQuotaManager.updateQuota(
- client1.configUser,
- client1.configClientEntity,
- Optional.of(new Quota(2000, true))
- )
- clientQuotaManager.updateQuota(
- client2.configUser,
- client2.configClientEntity,
- Optional.of(new Quota(4000, true))
- )
-
- assertEquals(Long.MaxValue.toDouble,
clientQuotaManager.quota(randomClient.user, randomClient.clientId).bound, 0.0,
- "Default producer quota should be " + Long.MaxValue.toDouble)
- assertEquals(2000, clientQuotaManager.quota(client1.user,
client1.clientId).bound, 0.0,
- "Should return the overridden value (2000)")
- assertEquals(4000, clientQuotaManager.quota(client2.user,
client2.clientId).bound, 0.0,
- "Should return the overridden value (4000)")
-
- // p1 should be throttled using the overridden quota
- var throttleTimeMs = maybeRecord(clientQuotaManager, client1.user,
client1.clientId, 2500 * config.numQuotaSamples)
- assertTrue(throttleTimeMs > 0, s"throttleTimeMs should be > 0. was
$throttleTimeMs")
-
- // Case 2: Change quota again. The quota should be updated within
KafkaMetrics as well since the sensor was created.
- // p1 should not longer be throttled after the quota change
- clientQuotaManager.updateQuota(
- client1.configUser,
- client1.configClientEntity,
- Optional.of(new Quota(3000, true))
- )
- assertEquals(3000, clientQuotaManager.quota(client1.user,
client1.clientId).bound, 0.0, "Should return the newly overridden value (3000)")
-
- throttleTimeMs = maybeRecord(clientQuotaManager, client1.user,
client1.clientId, 0)
- assertEquals(0, throttleTimeMs, s"throttleTimeMs should be 0. was
$throttleTimeMs")
-
- // Case 3: Change quota back to default. Should be throttled again
- clientQuotaManager.updateQuota(
- client1.configUser,
- client1.configClientEntity,
- Optional.of(new Quota(500, true))
- )
- assertEquals(500, clientQuotaManager.quota(client1.user,
client1.clientId).bound, 0.0, "Should return the default value (500)")
-
- throttleTimeMs = maybeRecord(clientQuotaManager, client1.user,
client1.clientId, 0)
- assertTrue(throttleTimeMs > 0, s"throttleTimeMs should be > 0. was
$throttleTimeMs")
-
- // Case 4: Set high default quota, remove p1 quota. p1 should no longer
be throttled
- clientQuotaManager.updateQuota(
- client1.configUser,
- client1.configClientEntity,
- Optional.empty
- )
- clientQuotaManager.updateQuota(
- defaultConfigClient.configUser,
- defaultConfigClient.configClientEntity,
- Optional.of(new Quota(4000, true))
- )
- assertEquals(4000, clientQuotaManager.quota(client1.user,
client1.clientId).bound, 0.0, "Should return the newly overridden value (4000)")
-
- throttleTimeMs = maybeRecord(clientQuotaManager, client1.user,
client1.clientId, 1000 * config.numQuotaSamples)
- assertEquals(0, throttleTimeMs, s"throttleTimeMs should be 0. was
$throttleTimeMs")
-
- } finally {
- clientQuotaManager.shutdown()
- }
- }
-
- /**
- * Tests parsing for <user> quotas when client-id default quota properties
are set.
- */
- @Test
- def testUserQuotaParsingWithDefaultClientIdQuota(): Unit = {
- val client1 = UserClient("User1", "p1", Optional.of(new
ClientQuotaManager.UserEntity("User1")), Optional.empty)
- val client2 = UserClient("User2", "p2", Optional.of(new
ClientQuotaManager.UserEntity("User2")), Optional.empty)
- val randomClient = UserClient("RandomUser", "random-client-id",
Optional.empty, Optional.empty)
- val defaultConfigClient = UserClient("", "",
Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY), Optional.empty)
- testQuotaParsing(config, client1, client2, randomClient,
defaultConfigClient)
- }
-
- private def checkQuota(quotaManager: ClientQuotaManager, user: String,
clientId: String, expectedBound: Long, value: Int, expectThrottle: Boolean):
Unit = {
- assertEquals(expectedBound.toDouble, quotaManager.quota(user,
clientId).bound, 0.0)
- val session = new Session(new KafkaPrincipal(KafkaPrincipal.USER_TYPE,
user), InetAddress.getLocalHost)
- val expectedMaxValueInQuotaWindow =
- if (expectedBound < Long.MaxValue) config.quotaWindowSizeSeconds *
(config.numQuotaSamples - 1) * expectedBound.toDouble
- else Double.MaxValue
- assertEquals(expectedMaxValueInQuotaWindow,
quotaManager.maxValueInQuotaWindow(session, clientId), 0.01)
-
- val throttleTimeMs = maybeRecord(quotaManager, user, clientId, value *
config.numQuotaSamples)
- if (expectThrottle)
- assertTrue(throttleTimeMs > 0, s"throttleTimeMs should be > 0. was
$throttleTimeMs")
- else
- assertEquals(0, throttleTimeMs, s"throttleTimeMs should be 0. was
$throttleTimeMs")
- }
-
- @Test
- def testMaxValueInQuotaWindowWithNonDefaultQuotaWindow(): Unit = {
- val numFullQuotaWindows = 3 // 3 seconds window (vs. 10 seconds default)
- val nonDefaultConfig = new ClientQuotaManagerConfig(numFullQuotaWindows +
1)
- val clientQuotaManager = new ClientQuotaManager(nonDefaultConfig, metrics,
QuotaType.FETCH, time, "")
- val userSession = new Session(new KafkaPrincipal(KafkaPrincipal.USER_TYPE,
"userA"), InetAddress.getLocalHost)
-
- try {
- // no quota set
- assertEquals(Double.MaxValue,
clientQuotaManager.maxValueInQuotaWindow(userSession, "client1"), 0.01)
-
- // Set default <user> quota config
- clientQuotaManager.updateQuota(
- Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY),
- Optional.empty,
- Optional.of(new Quota(10, true))
- )
- assertEquals(10 * numFullQuotaWindows,
clientQuotaManager.maxValueInQuotaWindow(userSession, "client1"), 0.01)
- } finally {
- clientQuotaManager.shutdown()
- }
- }
-
- @Test
- def testSetAndRemoveDefaultUserQuota(): Unit = {
- // quotaTypesEnabled will be QuotaTypes.NoQuotas initially
- val clientQuotaManager = new ClientQuotaManager(new
ClientQuotaManagerConfig(),
- metrics, QuotaType.PRODUCE, time, "")
-
- try {
- // no quota set yet, should not throttle
- checkQuota(clientQuotaManager, "userA", "client1", Long.MaxValue, 1000,
expectThrottle = false)
-
- // Set default <user> quota config
- clientQuotaManager.updateQuota(
- Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY),
- Optional.empty,
- Optional.of(new Quota(10, true))
- )
- checkQuota(clientQuotaManager, "userA", "client1", 10, 1000,
expectThrottle = true)
-
- // Remove default <user> quota config, back to no quotas
- clientQuotaManager.updateQuota(
- Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY),
- Optional.empty,
- Optional.empty
- )
- checkQuota(clientQuotaManager, "userA", "client1", Long.MaxValue, 1000,
expectThrottle = false)
- } finally {
- clientQuotaManager.shutdown()
- }
- }
-
- @Test
- def testSetAndRemoveUserQuota(): Unit = {
-
- val clientQuotaManager = new ClientQuotaManager(new
ClientQuotaManagerConfig(),
- metrics, QuotaType.PRODUCE, time, "")
-
- try {
- // Set <user> quota config
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userA")),
- Optional.empty,
- Optional.of(new Quota(10, true))
- )
- checkQuota(clientQuotaManager, "userA", "client1", 10, 1000,
expectThrottle = true)
-
- // Remove <user> quota config, back to no quotas
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userA")),
- Optional.empty,
- Optional.empty
- )
- checkQuota(clientQuotaManager, "userA", "client1", Long.MaxValue, 1000,
expectThrottle = false)
- } finally {
- clientQuotaManager.shutdown()
- }
- }
-
- @Test
- def testSetAndRemoveUserClientQuota(): Unit = {
- // quotaTypesEnabled will be QuotaTypes.NoQuotas initially
- val clientQuotaManager = new ClientQuotaManager(new
ClientQuotaManagerConfig(),
- metrics, QuotaType.PRODUCE, time, "")
-
- try {
- // Set <user, client-id> quota config
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userA")),
- Optional.of(new ClientQuotaManager.ClientIdEntity("client1")),
- Optional.of(new Quota(10, true))
- )
- checkQuota(clientQuotaManager, "userA", "client1", 10, 1000,
expectThrottle = true)
-
- // Remove <user, client-id> quota config, back to no quotas
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userA")),
- Optional.of(new ClientQuotaManager.ClientIdEntity("client1")),
- Optional.empty
- )
- checkQuota(clientQuotaManager, "userA", "client1", Long.MaxValue, 1000,
expectThrottle = false)
- } finally {
- clientQuotaManager.shutdown()
- }
- }
-
- @Test
- def testQuotaConfigPrecedence(): Unit = {
- val clientQuotaManager = new ClientQuotaManager(new
ClientQuotaManagerConfig(),
- metrics, QuotaType.PRODUCE, time, "")
-
- try {
- clientQuotaManager.updateQuota(
- Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY),
- Optional.empty,
- Optional.of(new Quota(1000, true))
- )
- clientQuotaManager.updateQuota(
- Optional.empty,
- Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
- Optional.of(new Quota(2000, true))
- )
- clientQuotaManager.updateQuota(
- Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY),
- Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
- Optional.of(new Quota(3000, true))
- )
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userA")),
- Optional.empty,
- Optional.of(new Quota(4000, true))
- )
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userA")),
- Optional.of(new ClientQuotaManager.ClientIdEntity("client1")),
- Optional.of(new Quota(5000, true))
- )
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userB")),
- Optional.empty,
- Optional.of(new Quota(6000, true))
- )
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userB")),
- Optional.of(new ClientQuotaManager.ClientIdEntity("client1")),
- Optional.of(new Quota(7000, true))
- )
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userB")),
- Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
- Optional.of(new Quota(8000, true))
- )
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userC")),
- Optional.empty,
- Optional.of(new Quota(10000, true))
- )
- clientQuotaManager.updateQuota(
- Optional.empty,
- Optional.of(new ClientQuotaManager.ClientIdEntity("client1")),
- Optional.of(new Quota(9000, true))
- )
-
- checkQuota(clientQuotaManager, "userA", "client1", 5000, 4500,
expectThrottle = false) // <user, client> quota takes precedence over <user>
- checkQuota(clientQuotaManager, "userA", "client2", 4000, 4500,
expectThrottle = true) // <user> quota takes precedence over <client> and
defaults
- checkQuota(clientQuotaManager, "userA", "client3", 4000, 0,
expectThrottle = true) // <user> quota is shared across clients of user
- checkQuota(clientQuotaManager, "userA", "client1", 5000, 0,
expectThrottle = false) // <user, client> is exclusive use, unaffected by
other clients
-
- checkQuota(clientQuotaManager, "userB", "client1", 7000, 8000,
expectThrottle = true)
- checkQuota(clientQuotaManager, "userB", "client2", 8000, 7000,
expectThrottle = false) // Default per-client quota for exclusive use of <user,
client>
- checkQuota(clientQuotaManager, "userB", "client3", 8000, 7000,
expectThrottle = false)
-
- checkQuota(clientQuotaManager, "userD", "client1", 3000, 3500,
expectThrottle = true) // Default <user, client> quota
- checkQuota(clientQuotaManager, "userD", "client2", 3000, 2500,
expectThrottle = false)
- checkQuota(clientQuotaManager, "userE", "client1", 3000, 2500,
expectThrottle = false)
-
- // Remove default <user, client> quota config, revert to <user> default
- clientQuotaManager.updateQuota(
- Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY),
- Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
- Optional.empty
- )
- checkQuota(clientQuotaManager, "userD", "client1", 1000, 0,
expectThrottle = false) // Metrics tags changed, restart counter
- checkQuota(clientQuotaManager, "userE", "client4", 1000, 1500,
expectThrottle = true)
- checkQuota(clientQuotaManager, "userF", "client4", 1000, 800,
expectThrottle = false) // Default <user> quota shared across clients of user
- checkQuota(clientQuotaManager, "userF", "client5", 1000, 800,
expectThrottle = true)
-
- // Remove default <user> quota config, revert to <client-id> default
- clientQuotaManager.updateQuota(
- Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY),
- Optional.empty,
- Optional.empty
- )
- checkQuota(clientQuotaManager, "userF", "client4", 2000, 0,
expectThrottle = false) // Default <client-id> quota shared across client-id
of all users
- checkQuota(clientQuotaManager, "userF", "client5", 2000, 0,
expectThrottle = false)
- checkQuota(clientQuotaManager, "userF", "client5", 2000, 2500,
expectThrottle = true)
- checkQuota(clientQuotaManager, "userG", "client5", 2000, 0,
expectThrottle = true)
-
- // Update quotas
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userA")),
- Optional.empty,
- Optional.of(new Quota(8000, true))
- )
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userA")),
- Optional.of(new ClientQuotaManager.ClientIdEntity("client1")),
- Optional.of(new Quota(10000, true))
- )
- checkQuota(clientQuotaManager, "userA", "client2", 8000, 0,
expectThrottle = false)
- checkQuota(clientQuotaManager, "userA", "client2", 8000, 4500,
expectThrottle = true) // Throttled due to sum of new and earlier values
- checkQuota(clientQuotaManager, "userA", "client1", 10000, 0,
expectThrottle = false)
- checkQuota(clientQuotaManager, "userA", "client1", 10000, 6000,
expectThrottle = true)
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userA")),
- Optional.of(new ClientQuotaManager.ClientIdEntity("client1")),
- Optional.empty
- )
- checkQuota(clientQuotaManager, "userA", "client6", 8000, 0,
expectThrottle = true) // Throttled due to shared user quota
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userA")),
- Optional.of(new ClientQuotaManager.ClientIdEntity("client6")),
- Optional.of(new Quota(11000, true))
- )
- checkQuota(clientQuotaManager, "userA", "client6", 11000, 8500,
expectThrottle = false)
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userA")),
- Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
- Optional.of(new Quota(12000, true))
- )
- clientQuotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity("userA")),
- Optional.of(new ClientQuotaManager.ClientIdEntity("client6")),
- Optional.empty
- )
- checkQuota(clientQuotaManager, "userA", "client6", 12000, 4000,
expectThrottle = true) // Throttled due to sum of new and earlier values
-
- } finally {
- clientQuotaManager.shutdown()
- }
- }
-
- @Test
- def testQuotaViolation(): Unit = {
- val clientQuotaManager = new ClientQuotaManager(config, metrics,
QuotaType.PRODUCE, time, "")
- val queueSizeMetric =
metrics.metrics().get(metrics.metricName("queue-size", "Produce", ""))
- try {
- clientQuotaManager.updateQuota(
- Optional.empty,
- Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
- Optional.of(new Quota(500, true))
- )
-
- // We have 10 seconds windows. Make sure that there is no quota violation
- // if we produce under the quota
- for (_ <- 0 until 10) {
- assertEquals(0, maybeRecord(clientQuotaManager, "ANONYMOUS",
"unknown", 400))
- time.sleep(1000)
- }
- assertEquals(0, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
-
- // Create a spike.
- // 400*10 + 2000 + 300 = 6300/10.5 = 600 bytes per second.
- // (600 - quota)/quota*window-size = (600-500)/500*10.5 seconds = 2100
- // 10.5 seconds because the last window is half complete
- time.sleep(500)
- val throttleTime = maybeRecord(clientQuotaManager, "ANONYMOUS",
"unknown", 2300)
-
- assertEquals(2100, throttleTime, "Should be throttled")
- throttle(clientQuotaManager, "ANONYMOUS", "unknown", throttleTime,
callback)
- assertEquals(1, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
- // After a request is delayed, the callback cannot be triggered
immediately
- clientQuotaManager.processThrottledChannelReaperDoWork
- assertEquals(0, numCallbacks)
- time.sleep(throttleTime)
-
- // Callback can only be triggered after the delay time passes
- clientQuotaManager.processThrottledChannelReaperDoWork()
- assertEquals(0, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
- assertEquals(1, numCallbacks)
-
- // Could continue to see delays until the bursty sample disappears
- for (_ <- 0 until 10) {
- maybeRecord(clientQuotaManager, "ANONYMOUS", "unknown", 400)
- time.sleep(1000)
- }
-
- assertEquals(0, maybeRecord(clientQuotaManager, "ANONYMOUS", "unknown",
0),
- "Should be unthrottled since bursty sample has rolled over")
- } finally {
- clientQuotaManager.shutdown()
- }
- }
-
- @Test
- def testExpireThrottleTimeSensor(): Unit = {
- val clientQuotaManager = new ClientQuotaManager(config, metrics,
QuotaType.PRODUCE, time, "")
- try {
- clientQuotaManager.updateQuota(
- Optional.empty,
- Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
- Optional.of(new Quota(500, true))
- )
-
- maybeRecord(clientQuotaManager, "ANONYMOUS", "client1", 100)
- // remove the throttle time sensor
- metrics.removeSensor("ProduceThrottleTime-:client1")
- // should not throw an exception even if the throttle time sensor does
not exist.
- val throttleTime = maybeRecord(clientQuotaManager, "ANONYMOUS",
"client1", 10000)
- assertTrue(throttleTime > 0, "Should be throttled")
- // the sensor should get recreated
- val throttleTimeSensor =
metrics.getSensor("ProduceThrottleTime-:client1")
- assertNotNull(throttleTimeSensor, "Throttle time sensor should exist")
- assertNotNull(throttleTimeSensor, "Throttle time sensor should exist")
- } finally {
- clientQuotaManager.shutdown()
- }
- }
-
- @Test
- def testExpireQuotaSensors(): Unit = {
- val clientQuotaManager = new ClientQuotaManager(config, metrics,
QuotaType.PRODUCE, time, "")
- try {
- clientQuotaManager.updateQuota(
- Optional.empty,
- Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
- Optional.of(new Quota(500, true))
- )
-
- maybeRecord(clientQuotaManager, "ANONYMOUS", "client1", 100)
- // remove all the sensors
- metrics.removeSensor("ProduceThrottleTime-:client1")
- metrics.removeSensor("Produce-ANONYMOUS:client1")
- // should not throw an exception
- val throttleTime = maybeRecord(clientQuotaManager, "ANONYMOUS",
"client1", 10000)
- assertTrue(throttleTime > 0, "Should be throttled")
-
- // all the sensors should get recreated
- val throttleTimeSensor =
metrics.getSensor("ProduceThrottleTime-:client1")
- assertNotNull(throttleTimeSensor, "Throttle time sensor should exist")
-
- val byteRateSensor = metrics.getSensor("Produce-:client1")
- assertNotNull(byteRateSensor, "Byte rate sensor should exist")
- } finally {
- clientQuotaManager.shutdown()
- }
- }
-
- @Test
- def testClientIdNotSanitized(): Unit = {
- val clientQuotaManager = new ClientQuotaManager(config, metrics,
QuotaType.PRODUCE, time, "")
- val clientId = "client@#$%"
- try {
- clientQuotaManager.updateQuota(
- Optional.empty,
- Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
- Optional.of(new Quota(500, true))
- )
-
- maybeRecord(clientQuotaManager, "ANONYMOUS", clientId, 100)
-
- // The metrics should use the raw client ID, even if the reporters
internally sanitize them
- val throttleTimeSensor = metrics.getSensor("ProduceThrottleTime-:" +
clientId)
- assertNotNull(throttleTimeSensor, "Throttle time sensor should exist")
-
- val byteRateSensor = metrics.getSensor("Produce-:" + clientId)
- assertNotNull(byteRateSensor, "Byte rate sensor should exist")
- } finally {
- clientQuotaManager.shutdown()
- }
- }
-
- @Test
- def testQuotaTypesEnabledUpdatesWithDefaultCallback(): Unit = {
- val clientQuotaManager = new ClientQuotaManager(config, metrics,
QuotaType.CONTROLLER_MUTATION, time, "")
- try {
- assertEquals(ClientQuotaManager.NO_QUOTAS,
clientQuotaManager.quotaTypesEnabled())
- assertFalse(clientQuotaManager.quotasEnabled)
-
- clientQuotaManager.updateQuota(Optional.empty(), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.of(new Quota(5, true)))
- assertEquals(ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED,
clientQuotaManager.quotaTypesEnabled())
- assertTrue(clientQuotaManager.quotasEnabled)
-
- clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.empty(), Optional.of(new
Quota(5, true)))
- assertEquals(ClientQuotaManager.USER_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED,
clientQuotaManager.quotaTypesEnabled())
- assertTrue(clientQuotaManager.quotasEnabled)
-
- clientQuotaManager.updateQuota(Optional.empty(), Optional.of(new
ClientQuotaManager.ClientIdEntity("client2")), Optional.of(new Quota(5, true)))
- assertEquals(ClientQuotaManager.USER_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED,
clientQuotaManager.quotaTypesEnabled())
- assertTrue(clientQuotaManager.quotasEnabled)
-
- clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userB")), Optional.empty(), Optional.of(new
Quota(5, true)))
- assertEquals(ClientQuotaManager.USER_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED,
clientQuotaManager.quotaTypesEnabled())
- assertTrue(clientQuotaManager.quotasEnabled)
-
- clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.of(new Quota(10, true)))
- assertEquals(ClientQuotaManager.USER_CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.USER_QUOTA_ENABLED, clientQuotaManager.quotaTypesEnabled())
- assertTrue(clientQuotaManager.quotasEnabled)
-
- clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.of(new Quota(12, true)))
- assertEquals(ClientQuotaManager.USER_CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.USER_QUOTA_ENABLED, clientQuotaManager.quotaTypesEnabled)
- assertTrue(clientQuotaManager.quotasEnabled)
-
- clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.empty(), Optional.empty())
- assertEquals(ClientQuotaManager.USER_CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.USER_QUOTA_ENABLED, clientQuotaManager.quotaTypesEnabled)
- assertTrue(clientQuotaManager.quotasEnabled)
-
- clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userB")), Optional.empty(), Optional.empty())
- assertEquals(ClientQuotaManager.USER_CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED,
clientQuotaManager.quotaTypesEnabled)
- assertTrue(clientQuotaManager.quotasEnabled)
-
- clientQuotaManager.updateQuota(Optional.empty(), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.empty())
- assertEquals(ClientQuotaManager.USER_CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED,
clientQuotaManager.quotaTypesEnabled)
- assertTrue(clientQuotaManager.quotasEnabled)
-
- clientQuotaManager.updateQuota(Optional.empty(), Optional.of(new
ClientQuotaManager.ClientIdEntity("client2")), Optional.empty())
- assertEquals(ClientQuotaManager.USER_CLIENT_ID_QUOTA_ENABLED,
clientQuotaManager.quotaTypesEnabled)
- assertTrue(clientQuotaManager.quotasEnabled)
-
- clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.empty())
- assertEquals(ClientQuotaManager.NO_QUOTAS,
clientQuotaManager.quotaTypesEnabled)
- assertFalse(clientQuotaManager.quotasEnabled)
-
- clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.empty())
- assertEquals(ClientQuotaManager.NO_QUOTAS,
clientQuotaManager.quotaTypesEnabled)
- assertFalse(clientQuotaManager.quotasEnabled)
- } finally {
- clientQuotaManager.shutdown()
- }
- }
-
- @Test
- def testQuotaTypesEnabledUpdatesWithCustomCallback(): Unit = {
- val customQuotaCallback = new ClientQuotaCallback {
- val quotas = new HashMap[ClientQuotaEntity, Quota]()
- override def configure(configs: Map[String, _]): Unit = {}
-
- override def quotaMetricTags(quotaType: ClientQuotaType, principal:
KafkaPrincipal, clientId: String): Map[String, String] = Collections.emptyMap()
-
- override def quotaLimit(quotaType: ClientQuotaType, metricTags:
Map[String, String]): java.lang.Double = 1
-
- override def updateQuota(quotaType: ClientQuotaType, entity:
ClientQuotaEntity, newValue: Double): Unit = {
- quotas.put(entity.asInstanceOf[ClientQuotaManager.KafkaQuotaEntity],
new Quota(newValue.toLong, true))
- }
-
- override def removeQuota(quotaType: ClientQuotaType, entity:
ClientQuotaEntity): Unit = {
- quotas.remove(entity.asInstanceOf[ClientQuotaManager.KafkaQuotaEntity])
- }
-
- override def quotaResetRequired(quotaType: ClientQuotaType): Boolean =
false
-
- override def close(): Unit = {}
- }
- val clientQuotaManager = new ClientQuotaManager(
- new ClientQuotaManagerConfig(),
- metrics,
- QuotaType.CONTROLLER_MUTATION,
- time,
- "",
- Optional.of(Plugin.wrapInstance(customQuotaCallback, metrics, ""))
- )
-
- try {
- assertEquals(ClientQuotaManager.CUSTOM_QUOTAS,
clientQuotaManager.quotaTypesEnabled)
- assertTrue(clientQuotaManager.quotasEnabled, "quotasEnabled should be
true with custom callback")
-
- clientQuotaManager.updateQuota(Optional.empty(), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.of(new Quota(12, true)))
- assertEquals(ClientQuotaManager.CUSTOM_QUOTAS,
clientQuotaManager.quotaTypesEnabled)
-
- clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.empty(), Optional.of(new
Quota(12, true)))
- assertEquals(ClientQuotaManager.CUSTOM_QUOTAS,
clientQuotaManager.quotaTypesEnabled)
- assertTrue(clientQuotaManager.quotasEnabled, "quotasEnabled should
remain true")
-
- clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.of(new Quota(12, true)))
- assertEquals(ClientQuotaManager.CUSTOM_QUOTAS,
clientQuotaManager.quotaTypesEnabled())
- assertTrue(clientQuotaManager.quotasEnabled, "quotasEnabled should
remain true")
-
- clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.empty())
- clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.empty(), Optional.empty())
- clientQuotaManager.updateQuota(Optional.empty(), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.empty())
- assertEquals(ClientQuotaManager.CUSTOM_QUOTAS,
clientQuotaManager.quotaTypesEnabled())
- assertTrue(clientQuotaManager.quotasEnabled, "quotasEnabled should
remain true")
- } finally {
- clientQuotaManager.shutdown()
- }
- }
-
- private case class UserClient(
- user: String,
- clientId: String,
- configUser: Optional[ClientQuotaEntity.ConfigEntity] = Optional.empty,
- configClientEntity: Optional[ClientQuotaEntity.ConfigEntity] =
Optional.empty
- )
-}
diff --git
a/core/src/test/scala/unit/kafka/server/ClientRequestQuotaManagerTest.scala
b/core/src/test/scala/unit/kafka/server/ClientRequestQuotaManagerTest.scala
deleted file mode 100644
index 9b8e85c44e7..00000000000
--- a/core/src/test/scala/unit/kafka/server/ClientRequestQuotaManagerTest.scala
+++ /dev/null
@@ -1,99 +0,0 @@
-/**
- * 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 org.apache.kafka.common.metrics.Quota
-import org.apache.kafka.server.config.ClientQuotaManagerConfig
-import org.apache.kafka.server.quota.{ClientQuotaManager, QuotaType}
-import org.apache.kafka.server.quota.ClientQuotaEntity
-import org.junit.jupiter.api.Assertions._
-import org.junit.jupiter.api.Test
-
-import java.util.Optional
-
-class ClientRequestQuotaManagerTest extends BaseClientQuotaManagerTest {
- private val config = new ClientQuotaManagerConfig()
-
- @Test
- def testRequestPercentageQuotaViolation(): Unit = {
- val clientRequestQuotaManager = new ClientRequestQuotaManager(config,
metrics, time, "", Optional.empty())
- val userEntity: ClientQuotaEntity.ConfigEntity = new
ClientQuotaManager.UserEntity("ANONYMOUS")
- val clientEntity: ClientQuotaEntity.ConfigEntity = new
ClientQuotaManager.ClientIdEntity("test-client")
-
- clientRequestQuotaManager.updateQuota(
- Optional.of(userEntity),
- Optional.of(clientEntity),
- Optional.of(Quota.upperBound(1))
- )
- val queueSizeMetric =
metrics.metrics().get(metrics.metricName("queue-size",
QuotaType.REQUEST.toString, ""))
- def millisToPercent(millis: Double) = millis * 1000 * 1000 *
ClientRequestQuotaManager.NANOS_TO_PERCENTAGE_PER_SECOND
- try {
- // We have 10 seconds windows. Make sure that there is no quota violation
- // if we are under the quota
- for (_ <- 0 until 10) {
- assertEquals(0, maybeRecord(clientRequestQuotaManager, "ANONYMOUS",
"test-client", millisToPercent(4)))
- time.sleep(1000)
- }
- assertEquals(0, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
-
- // Create a spike.
- // quota = 1% (10ms per second)
- // 4*10 + 67.1 = 107.1/10.5 = 10.2ms per second.
- // (10.2 - quota)/quota*window-size = (10.2-10)/10*10.5 seconds = 210ms
- // 10.5 seconds interval because the last window is half complete
- time.sleep(500)
- val throttleTime = maybeRecord(clientRequestQuotaManager, "ANONYMOUS",
"test-client", millisToPercent(67.1))
-
- assertEquals(210, throttleTime, "Should be throttled")
-
- throttle(clientRequestQuotaManager, "ANONYMOUS", "test-client",
throttleTime, callback)
- assertEquals(1, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
- // After a request is delayed, the callback cannot be triggered
immediately
- clientRequestQuotaManager.processThrottledChannelReaperDoWork()
- assertEquals(0, numCallbacks)
- time.sleep(throttleTime)
-
- // Callback can only be triggered after the delay time passes
- clientRequestQuotaManager.processThrottledChannelReaperDoWork()
- assertEquals(0, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
- assertEquals(1, numCallbacks)
-
- // Could continue to see delays until the bursty sample disappears
- for (_ <- 0 until 11) {
- maybeRecord(clientRequestQuotaManager, "ANONYMOUS", "test-client",
millisToPercent(4))
- time.sleep(1000)
- }
-
- assertEquals(0,
- maybeRecord(clientRequestQuotaManager, "ANONYMOUS", "test-client", 0),
"Should be unthrottled since bursty sample has rolled over")
-
- // Create a very large spike which requires > one quota window to bring
within quota
- assertEquals(1000, maybeRecord(clientRequestQuotaManager, "ANONYMOUS",
"test-client", millisToPercent(500)))
- for (_ <- 0 until 10) {
- time.sleep(1000)
- assertEquals(1000, maybeRecord(clientRequestQuotaManager, "ANONYMOUS",
"test-client", 0))
- }
- time.sleep(1000)
- assertEquals(0,
- maybeRecord(clientRequestQuotaManager, "ANONYMOUS", "test-client", 0),
"Should be unthrottled since bursty sample has rolled over")
-
- } finally {
- clientRequestQuotaManager.shutdown()
- }
- }
-}
-
diff --git a/core/src/test/scala/unit/kafka/server/ControllerApisTest.scala
b/core/src/test/scala/unit/kafka/server/ControllerApisTest.scala
index fd189d1c95e..da1d4d2112a 100644
--- a/core/src/test/scala/unit/kafka/server/ControllerApisTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ControllerApisTest.scala
@@ -59,7 +59,7 @@ import org.apache.kafka.server.SimpleApiVersionManager
import org.apache.kafka.server.authorizer.{Action, AuthorizableRequestContext,
AuthorizationResult, Authorizer}
import org.apache.kafka.server.common.{ApiMessageAndVersion,
FinalizedFeatures, KRaftVersion, MetadataVersion, ProducerIdsBlock,
RequestLocal}
import org.apache.kafka.server.config.ServerConfigs
-import org.apache.kafka.server.quota.{ClientQuotaManager,
ControllerMutationQuota, ControllerMutationQuotaManager,
ReplicationQuotaManager}
+import org.apache.kafka.server.quota.{ClientQuotaManager,
ClientRequestQuotaManager, ControllerMutationQuota,
ControllerMutationQuotaManager, ReplicationQuotaManager}
import org.apache.kafka.storage.internals.log.CleanerConfig
import org.apache.kafka.test.TestUtils
import org.junit.jupiter.api.Assertions._
diff --git
a/core/src/test/scala/unit/kafka/server/ControllerMutationQuotaManagerTest.scala
b/core/src/test/scala/unit/kafka/server/ControllerMutationQuotaManagerTest.scala
deleted file mode 100644
index a40087a5973..00000000000
---
a/core/src/test/scala/unit/kafka/server/ControllerMutationQuotaManagerTest.scala
+++ /dev/null
@@ -1,243 +0,0 @@
-/**
- * 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 java.util.concurrent.TimeUnit
-import org.apache.kafka.common.errors.ThrottlingQuotaExceededException
-import org.apache.kafka.common.metrics.MetricConfig
-import org.apache.kafka.common.metrics.Metrics
-import org.apache.kafka.common.metrics.Quota
-import org.apache.kafka.common.metrics.QuotaViolationException
-import org.apache.kafka.common.metrics.stats.TokenBucket
-import org.apache.kafka.common.utils.MockTime
-import org.apache.kafka.server.config.ClientQuotaManagerConfig
-import org.apache.kafka.server.quota.{ClientQuotaManager,
ControllerMutationQuota, ControllerMutationQuotaManager,
PermissiveControllerMutationQuota, QuotaType, StrictControllerMutationQuota}
-import org.junit.jupiter.api.Assertions._
-import org.junit.jupiter.api.Assertions.assertEquals
-import org.junit.jupiter.api.Assertions.assertFalse
-import org.junit.jupiter.api.Test
-
-import java.util.Optional
-
-class StrictControllerMutationQuotaTest {
- @Test
- def testControllerMutationQuotaViolation(): Unit = {
- val time = new MockTime(0, System.currentTimeMillis, 0)
- val metrics = new Metrics(time)
- val sensor = metrics.sensor("sensor", new MetricConfig()
- .quota(Quota.upperBound(10))
- .timeWindow(1, TimeUnit.SECONDS)
- .samples(10))
- val metricName = metrics.metricName("rate", "test-group")
- assertTrue(sensor.add(metricName, new TokenBucket))
-
- val quota = new StrictControllerMutationQuota(time, sensor)
- assertFalse(quota.isExceeded)
-
- // Recording a first value at T to bring the tokens to 10. Value is
accepted
- // because the quota is not exhausted yet.
- quota.record(90)
- assertFalse(quota.isExceeded)
- assertEquals(0, quota.throttleTime)
-
- // Recording a second value at T to bring the tokens to -80. Value is
accepted
- quota.record(90)
- assertFalse(quota.isExceeded)
- assertEquals(0, quota.throttleTime)
-
- // Recording a third value at T is rejected immediately because there are
not
- // tokens available in the bucket.
- assertThrows(classOf[ThrottlingQuotaExceededException], () =>
quota.record(90))
- assertTrue(quota.isExceeded)
- assertEquals(8000, quota.throttleTime)
-
- // Throttle time is adjusted with time
- time.sleep(5000)
- assertEquals(3000, quota.throttleTime)
-
- metrics.close()
- }
-}
-
-class PermissiveControllerMutationQuotaTest {
- @Test
- def testControllerMutationQuotaViolation(): Unit = {
- val time = new MockTime(0, System.currentTimeMillis, 0)
- val metrics = new Metrics(time)
- val sensor = metrics.sensor("sensor", new MetricConfig()
- .quota(Quota.upperBound(10))
- .timeWindow(1, TimeUnit.SECONDS)
- .samples(10))
- val metricName = metrics.metricName("rate", "test-group")
- assertTrue(sensor.add(metricName, new TokenBucket))
-
- val quota = new PermissiveControllerMutationQuota(time, sensor)
- assertFalse(quota.isExceeded)
-
- // Recording a first value at T to bring the tokens 10. Value is accepted
- // because the quota is not exhausted yet.
- quota.record(90)
- assertFalse(quota.isExceeded)
- assertEquals(0, quota.throttleTime)
-
- // Recording a second value at T to bring the tokens to -80. Value is
accepted
- quota.record(90)
- assertFalse(quota.isExceeded)
- assertEquals(8000, quota.throttleTime)
-
- // Recording a second value at T to bring the tokens to -170. Value is
accepted
- // even though the quota is exhausted.
- quota.record(90)
- assertFalse(quota.isExceeded) // quota is never exceeded
- assertEquals(17000, quota.throttleTime)
-
- // Throttle time is adjusted with time
- time.sleep(5000)
- assertEquals(12000, quota.throttleTime)
-
- metrics.close()
- }
-}
-
-class ControllerMutationQuotaManagerTest extends BaseClientQuotaManagerTest {
- private val User = "ANONYMOUS"
- private val ClientId = "test-client"
-
- private val config = new ClientQuotaManagerConfig(10, 1)
-
- private def withQuotaManager(f: ControllerMutationQuotaManager => Unit):
Unit = {
- val quotaManager = new ControllerMutationQuotaManager(config, metrics,
time,"", Optional.empty())
- try {
- f(quotaManager)
- } finally {
- quotaManager.shutdown()
- }
- }
-
- @Test
- def testThrottleTime(): Unit = {
- import ControllerMutationQuotaManager._
-
- val time = new MockTime(0, System.currentTimeMillis, 0)
- val metrics = new Metrics(time)
- val sensor = metrics.sensor("sensor")
- val metricName = metrics.metricName("tokens", "test-group")
- sensor.add(metricName, new TokenBucket)
- val metric = metrics.metric(metricName)
-
- assertEquals(0, throttleTimeMs(new QuotaViolationException(metric, 0, 10)))
- assertEquals(500, throttleTimeMs(new QuotaViolationException(metric, -5,
10)))
- assertEquals(1000, throttleTimeMs(new QuotaViolationException(metric, -10,
10)))
- }
-
- @Test
- def testControllerMutationQuotaViolation(): Unit = {
- withQuotaManager { quotaManager =>
- quotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity(User)),
- Optional.of(new ClientQuotaManager.ClientIdEntity(ClientId)),
- Optional.of(Quota.upperBound(10))
- )
- val queueSizeMetric = metrics.metrics().get(
- metrics.metricName("queue-size",
QuotaType.CONTROLLER_MUTATION.toString, ""))
-
- // Verify that there is no quota violation if we remain under the quota.
- for (_ <- 0 until 10) {
- assertEquals(0, maybeRecord(quotaManager, User, ClientId, 10))
- time.sleep(1000)
- }
- assertEquals(0, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
-
- // Create a spike worth of 110 mutations.
- // Current tokens in the bucket = 100
- // As we use the Strict enforcement, the quota is checked before
updating the rate. Hence,
- // the spike is accepted and no quota violation error is raised.
- var throttleTime = maybeRecord(quotaManager, User, ClientId, 110)
- assertEquals(0, throttleTime, "Should not be throttled")
-
- // Create a spike worth of 110 mutations.
- // Current tokens in the bucket = 100 - 110 = -10
- // As the quota is already violated, the spike is rejected immediately
without updating the
- // rate. The client must wait:
- // 10 / 10 = 1s
- throttleTime = maybeRecord(quotaManager, User, ClientId, 110)
- assertEquals(1000, throttleTime, "Should be throttled")
-
- // Throttle
- throttle(quotaManager, User, ClientId, throttleTime, callback)
- assertEquals(1, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
-
- // After a request is delayed, the callback cannot be triggered
immediately
- quotaManager.processThrottledChannelReaperDoWork()
- assertEquals(0, numCallbacks)
-
- // Callback can only be triggered after the delay time passes
- time.sleep(throttleTime)
- quotaManager.processThrottledChannelReaperDoWork()
- assertEquals(0, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
- assertEquals(1, numCallbacks)
-
- // Retry to spike worth of 110 mutations after having waited the
required throttle time.
- // Current tokens in the bucket = 0
- throttleTime = maybeRecord(quotaManager, User, ClientId, 110)
- assertEquals(0, throttleTime, "Should be throttled")
- }
- }
-
- @Test
- def testNewStrictQuotaForReturnsUnboundedQuotaWhenQuotaIsDisabled(): Unit = {
- withQuotaManager { quotaManager =>
- assertEquals(ControllerMutationQuota.UNBOUNDED_CONTROLLER_MUTATION_QUOTA,
- quotaManager.newStrictQuotaFor(buildSession(User), ClientId))
- }
- }
-
- @Test
- def testNewStrictQuotaForReturnsStrictQuotaWhenQuotaIsEnabled(): Unit = {
- withQuotaManager { quotaManager =>
- quotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity(User)),
- Optional.of(new ClientQuotaManager.ClientIdEntity(ClientId)),
- Optional.of(Quota.upperBound(10))
- )
- val quota = quotaManager.newStrictQuotaFor(buildSession(User), ClientId)
- assertTrue(quota.isInstanceOf[StrictControllerMutationQuota])
-
- }
- }
-
- @Test
- def testNewPermissiveQuotaForReturnsUnboundedQuotaWhenQuotaIsDisabled():
Unit = {
- withQuotaManager { quotaManager =>
- assertEquals(ControllerMutationQuota.UNBOUNDED_CONTROLLER_MUTATION_QUOTA,
- quotaManager.newPermissiveQuotaFor(buildSession(User), ClientId))
- }
- }
-
- @Test
- def testNewPermissiveQuotaForReturnsStrictQuotaWhenQuotaIsEnabled(): Unit = {
- withQuotaManager { quotaManager =>
- quotaManager.updateQuota(
- Optional.of(new ClientQuotaManager.UserEntity(User)),
- Optional.of(new ClientQuotaManager.ClientIdEntity(ClientId)),
- Optional.of(Quota.upperBound(10))
- )
- val quota = quotaManager.newPermissiveQuotaFor(buildSession(User),
ClientId)
- assertTrue(quota.isInstanceOf[PermissiveControllerMutationQuota])
- }
- }
-}
diff --git a/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
b/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
index 6310d0a2f6c..f5d838edf83 100644
--- a/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
@@ -99,7 +99,7 @@ import org.apache.kafka.server.config.{ReplicationConfigs,
ServerConfigs, Server
import org.apache.kafka.server.logger.LoggingController
import org.apache.kafka.server.metrics.{ClientMetricsTestUtils,
KafkaYammerMetrics}
import org.apache.kafka.server.share.{CachedSharePartition,
ErroneousAndValidPartitionData, SharePartitionKey}
-import org.apache.kafka.server.quota.{ClientQuotaManager,
ControllerMutationQuota, ControllerMutationQuotaManager, ReplicaQuota,
ReplicationQuotaManager, ThrottleCallback}
+import org.apache.kafka.server.quota.{ClientQuotaManager,
ClientRequestQuotaManager, ControllerMutationQuota,
ControllerMutationQuotaManager, ReplicaQuota, ReplicationQuotaManager,
ThrottleCallback}
import org.apache.kafka.server.share.acknowledge.ShareAcknowledgementBatch
import org.apache.kafka.server.share.context.{FinalContext,
ShareSessionContext}
import org.apache.kafka.server.share.session.{ShareSession, ShareSessionKey}
diff --git
a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/metadata/KRaftMetadataRequestBenchmark.java
b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/metadata/KRaftMetadataRequestBenchmark.java
index 227951f075c..0dad41f0f87 100644
---
a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/metadata/KRaftMetadataRequestBenchmark.java
+++
b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/metadata/KRaftMetadataRequestBenchmark.java
@@ -19,7 +19,6 @@ package org.apache.kafka.jmh.metadata;
import kafka.coordinator.transaction.TransactionCoordinator;
import kafka.network.RequestChannel;
-import kafka.server.ClientRequestQuotaManager;
import kafka.server.KafkaApis;
import kafka.server.KafkaConfig;
import kafka.server.QuotaFactory;
@@ -64,6 +63,7 @@ import org.apache.kafka.server.common.FinalizedFeatures;
import org.apache.kafka.server.common.KRaftVersion;
import org.apache.kafka.server.common.MetadataVersion;
import org.apache.kafka.server.quota.ClientQuotaManager;
+import org.apache.kafka.server.quota.ClientRequestQuotaManager;
import org.apache.kafka.server.quota.ControllerMutationQuotaManager;
import org.apache.kafka.server.quota.ReplicationQuotaManager;
import org.apache.kafka.storage.log.metrics.BrokerTopicStats;
diff --git a/core/src/main/java/kafka/server/ClientRequestQuotaManager.java
b/server/src/main/java/org/apache/kafka/server/quota/ClientRequestQuotaManager.java
similarity index 95%
rename from core/src/main/java/kafka/server/ClientRequestQuotaManager.java
rename to
server/src/main/java/org/apache/kafka/server/quota/ClientRequestQuotaManager.java
index e685b6aa50a..b28ba68d9ab 100644
--- a/core/src/main/java/kafka/server/ClientRequestQuotaManager.java
+++
b/server/src/main/java/org/apache/kafka/server/quota/ClientRequestQuotaManager.java
@@ -14,7 +14,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package kafka.server;
+package org.apache.kafka.server.quota;
import org.apache.kafka.common.MetricName;
import org.apache.kafka.common.internals.Plugin;
@@ -25,10 +25,6 @@ import org.apache.kafka.common.metrics.stats.Rate;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.network.Request;
import org.apache.kafka.server.config.ClientQuotaManagerConfig;
-import org.apache.kafka.server.quota.ClientQuotaCallback;
-import org.apache.kafka.server.quota.ClientQuotaManager;
-import org.apache.kafka.server.quota.QuotaType;
-import org.apache.kafka.server.quota.QuotaUtils;
import java.util.Map;
import java.util.Optional;
diff --git
a/server/src/test/java/org/apache/kafka/server/quota/BaseClientQuotaManagerTest.java
b/server/src/test/java/org/apache/kafka/server/quota/BaseClientQuotaManagerTest.java
new file mode 100644
index 00000000000..c01ed490a67
--- /dev/null
+++
b/server/src/test/java/org/apache/kafka/server/quota/BaseClientQuotaManagerTest.java
@@ -0,0 +1,97 @@
+/*
+ * 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 org.apache.kafka.server.quota;
+
+import org.apache.kafka.common.memory.MemoryPool;
+import org.apache.kafka.common.metrics.MetricConfig;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.network.ClientInformation;
+import org.apache.kafka.common.network.ListenerName;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.requests.AbstractRequest;
+import org.apache.kafka.common.requests.FetchRequest;
+import org.apache.kafka.common.requests.RequestContext;
+import org.apache.kafka.common.requests.RequestHeader;
+import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.network.Request;
+import org.apache.kafka.network.Session;
+import org.apache.kafka.network.metrics.RequestChannelMetrics;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+
+import java.net.InetAddress;
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.mockito.Mockito.mock;
+
+public class BaseClientQuotaManagerTest {
+
+ protected final MockTime time = new MockTime();
+ protected int numCallbacks = 0;
+ protected Metrics metrics;
+
+ @BeforeEach
+ public void setup() {
+ metrics = new Metrics(new MetricConfig(), List.of(), time);
+ }
+
+ @AfterEach
+ public void tearDown() {
+ metrics.close();
+ }
+
+ protected final ThrottleCallback callback = new ThrottleCallback() {
+ @Override
+ public void startThrottling() {}
+
+ @Override
+ public void endThrottling() {
+ // Count how many times this callback is called for
notifyThrottlingDone().
+ numCallbacks += 1;
+ }
+ };
+
+ protected Session buildSession(String user) {
+ KafkaPrincipal principal = new
KafkaPrincipal(KafkaPrincipal.USER_TYPE, user);
+ return new Session(principal, null);
+ }
+
+ protected int maybeRecord(ClientQuotaManager quotaManager, String user,
String clientId, double value) {
+ return
quotaManager.maybeRecordAndGetThrottleTimeMs(buildSession(user), clientId,
value, time.milliseconds());
+ }
+
+ protected void throttle(ClientQuotaManager quotaManager, int
throttleTimeMs, ThrottleCallback channelThrottlingCallback) {
+ AbstractRequest.Builder<FetchRequest> builder =
FetchRequest.Builder.forConsumer(ApiKeys.FETCH.latestVersion(), 0, 1000,
Map.of());
+ FetchRequest fetchRequest = builder.build();
+ ByteBuffer buffer = fetchRequest.serializeWithHeader(new
RequestHeader(builder.apiKey(), fetchRequest.version(), "", 0));
+ RequestChannelMetrics requestChannelMetrics =
mock(RequestChannelMetrics.class);
+
+ // read the header from the buffer first so that the body can be read
next from the Request constructor
+ RequestHeader header = RequestHeader.parse(buffer);
+ RequestContext context = new RequestContext(header, "1",
assertDoesNotThrow(InetAddress::getLocalHost), KafkaPrincipal.ANONYMOUS,
+ ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
SecurityProtocol.PLAINTEXT, ClientInformation.EMPTY, false);
+ Request request = new Request(1, context, 0, MemoryPool.NONE, buffer,
requestChannelMetrics);
+ quotaManager.throttle(request.header().clientId(), request.session(),
channelThrottlingCallback, throttleTimeMs);
+ }
+
+}
diff --git
a/server/src/test/java/org/apache/kafka/server/quota/ClientQuotaManagerTest.java
b/server/src/test/java/org/apache/kafka/server/quota/ClientQuotaManagerTest.java
new file mode 100644
index 00000000000..ac841e47ec9
--- /dev/null
+++
b/server/src/test/java/org/apache/kafka/server/quota/ClientQuotaManagerTest.java
@@ -0,0 +1,620 @@
+/*
+ * 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 org.apache.kafka.server.quota;
+
+import org.apache.kafka.common.internals.Plugin;
+import org.apache.kafka.common.metrics.KafkaMetric;
+import org.apache.kafka.common.metrics.Quota;
+import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.apache.kafka.network.Session;
+import org.apache.kafka.server.config.ClientQuotaManagerConfig;
+
+import org.junit.jupiter.api.Test;
+
+import java.net.InetAddress;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.function.Consumer;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
+
+ private final ClientQuotaManagerConfig config = new
ClientQuotaManagerConfig();
+
+ private void withQuotaManager(QuotaType quotaType,
ClientQuotaManagerConfig config, Consumer<ClientQuotaManager> f) {
+ ClientQuotaManager quotaManager = new ClientQuotaManager(config,
metrics, quotaType, time, "");
+ try {
+ f.accept(quotaManager);
+ } finally {
+ quotaManager.shutdown();
+ }
+ }
+
+ private void withQuotaManager(QuotaType quotaType,
Consumer<ClientQuotaManager> f) {
+ withQuotaManager(quotaType, config, f);
+ }
+
+ private void testQuotaParsing(ClientQuotaManagerConfig config, UserClient
client1, UserClient client2, UserClient randomClient, UserClient
defaultConfigClient) {
+ withQuotaManager(QuotaType.PRODUCE, clientQuotaManager -> {
+ // Case 1: Update the quota. Assert that the new quota value is
returned
+ clientQuotaManager.updateQuota(
+ client1.configUser,
+ client1.configClientEntity,
+ Optional.of(new Quota(2000, true))
+ );
+ clientQuotaManager.updateQuota(
+ client2.configUser,
+ client2.configClientEntity,
+ Optional.of(new Quota(4000, true))
+ );
+
+ assertEquals(Long.MAX_VALUE,
clientQuotaManager.quota(randomClient.user, randomClient.clientId).bound(), 0.0,
+ "Default producer quota should be " + Long.MAX_VALUE);
+ assertEquals(2000, clientQuotaManager.quota(client1.user,
client1.clientId).bound(), 0.0,
+ "Should return the overridden value (2000)");
+ assertEquals(4000, clientQuotaManager.quota(client2.user,
client2.clientId).bound(), 0.0,
+ "Should return the overridden value (4000)");
+
+ // p1 should be throttled using the overridden quota
+ int throttleTimeMs = maybeRecord(clientQuotaManager, client1.user,
client1.clientId, 2500 * config.numQuotaSamples());
+ assertTrue(throttleTimeMs > 0, "throttleTimeMs should be > 0. was
" + throttleTimeMs);
+
+ // Case 2: Change quota again. The quota should be updated within
KafkaMetrics as well since the sensor was created.
+ // p1 should no longer be throttled after the quota change
+ clientQuotaManager.updateQuota(
+ client1.configUser,
+ client1.configClientEntity,
+ Optional.of(new Quota(3000, true))
+ );
+ assertEquals(3000, clientQuotaManager.quota(client1.user,
client1.clientId).bound(), 0.0, "Should return the newly overridden value
(3000)");
+
+ throttleTimeMs = maybeRecord(clientQuotaManager, client1.user,
client1.clientId, 0);
+ assertEquals(0, throttleTimeMs, "throttleTimeMs should be 0. was "
+ throttleTimeMs);
+
+ // Case 3: Change quota back to default. Should be throttled again
+ clientQuotaManager.updateQuota(
+ client1.configUser,
+ client1.configClientEntity,
+ Optional.of(new Quota(500, true))
+ );
+ assertEquals(500, clientQuotaManager.quota(client1.user,
client1.clientId).bound(), 0.0, "Should return the default value (500)");
+
+ throttleTimeMs = maybeRecord(clientQuotaManager, client1.user,
client1.clientId, 0);
+ assertTrue(throttleTimeMs > 0, "throttleTimeMs should be > 0. was
" + throttleTimeMs);
+
+ // Case 4: Set high default quota, remove p1 quota. p1 should no
longer be throttled
+ clientQuotaManager.updateQuota(
+ client1.configUser,
+ client1.configClientEntity,
+ Optional.empty()
+ );
+ clientQuotaManager.updateQuota(
+ defaultConfigClient.configUser,
+ defaultConfigClient.configClientEntity,
+ Optional.of(new Quota(4000, true))
+ );
+ assertEquals(4000, clientQuotaManager.quota(client1.user,
client1.clientId).bound(), 0.0, "Should return the newly overridden value
(4000)");
+
+ throttleTimeMs = maybeRecord(clientQuotaManager, client1.user,
client1.clientId, 1000 * config.numQuotaSamples());
+ assertEquals(0, throttleTimeMs, "throttleTimeMs should be 0. was "
+ throttleTimeMs);
+
+ });
+ }
+
+ /**
+ * Tests parsing for <user> quotas when client-id default quota properties
are set.
+ */
+ @Test
+ public void testUserQuotaParsingWithDefaultClientIdQuota() {
+ UserClient client1 = new UserClient("User1", "p1", Optional.of(new
ClientQuotaManager.UserEntity("User1")), Optional.empty());
+ UserClient client2 = new UserClient("User2", "p2", Optional.of(new
ClientQuotaManager.UserEntity("User2")), Optional.empty());
+ UserClient randomClient = new UserClient("RandomUser",
"random-client-id", Optional.empty(), Optional.empty());
+ UserClient defaultConfigClient = new UserClient("", "",
Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY), Optional.empty());
+ testQuotaParsing(config, client1, client2, randomClient,
defaultConfigClient);
+ }
+
+ private void checkQuota(ClientQuotaManager quotaManager, String user,
String clientId, long expectedBound, int value, boolean expectThrottle) {
+ assertEquals(expectedBound, quotaManager.quota(user,
clientId).bound(), 0.0);
+ Session session = new Session(new
KafkaPrincipal(KafkaPrincipal.USER_TYPE, user),
assertDoesNotThrow(InetAddress::getLocalHost));
+ double expectedMaxValueInQuotaWindow = expectedBound < Long.MAX_VALUE
+ ? config.quotaWindowSizeSeconds() * (config.numQuotaSamples() - 1)
* expectedBound
+ : Double.MAX_VALUE;
+ assertEquals(expectedMaxValueInQuotaWindow,
quotaManager.maxValueInQuotaWindow(session, clientId), 0.01);
+
+ int throttleTimeMs = maybeRecord(quotaManager, user, clientId, value *
config.numQuotaSamples());
+ if (expectThrottle) {
+ assertTrue(throttleTimeMs > 0, "throttleTimeMs should be > 0. was
" + throttleTimeMs);
+ } else {
+ assertEquals(0, throttleTimeMs, "throttleTimeMs should be 0. was "
+ throttleTimeMs);
+ }
+ }
+
+ @Test
+ public void testMaxValueInQuotaWindowWithNonDefaultQuotaWindow() {
+ int numFullQuotaWindows = 3; // 3 seconds window (vs. 10 seconds
default)
+ ClientQuotaManagerConfig nonDefaultConfig = new
ClientQuotaManagerConfig(numFullQuotaWindows + 1);
+ withQuotaManager(QuotaType.FETCH, nonDefaultConfig, clientQuotaManager
-> {
+ Session userSession = new Session(new
KafkaPrincipal(KafkaPrincipal.USER_TYPE, "userA"),
assertDoesNotThrow(InetAddress::getLocalHost));
+
+ // no quota set
+ assertEquals(Double.MAX_VALUE,
clientQuotaManager.maxValueInQuotaWindow(userSession, "client1"), 0.01);
+
+ // Set default <user> quota config
+ clientQuotaManager.updateQuota(
+ Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY),
+ Optional.empty(),
+ Optional.of(new Quota(10, true))
+ );
+ assertEquals(10 * numFullQuotaWindows,
clientQuotaManager.maxValueInQuotaWindow(userSession, "client1"), 0.01);
+ });
+ }
+
+ @Test
+ public void testSetAndRemoveDefaultUserQuota() {
+ // quotaTypesEnabled will be QuotaTypes.NoQuotas initially
+ withQuotaManager(QuotaType.PRODUCE, clientQuotaManager -> {
+ // no quota set yet, should not throttle
+ checkQuota(clientQuotaManager, "userA", "client1", Long.MAX_VALUE,
1000, false);
+
+ // Set default <user> quota config
+ clientQuotaManager.updateQuota(
+ Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY),
+ Optional.empty(),
+ Optional.of(new Quota(10, true))
+ );
+ checkQuota(clientQuotaManager, "userA", "client1", 10, 1000, true);
+
+ // Remove default <user> quota config, back to no quotas
+ clientQuotaManager.updateQuota(
+ Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY),
+ Optional.empty(),
+ Optional.empty()
+ );
+ checkQuota(clientQuotaManager, "userA", "client1", Long.MAX_VALUE,
1000, false);
+ });
+ }
+
+ @Test
+ public void testSetAndRemoveUserQuota() {
+ withQuotaManager(QuotaType.PRODUCE, clientQuotaManager -> {
+ // Set <user> quota config
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userA")),
+ Optional.empty(),
+ Optional.of(new Quota(10, true))
+ );
+ checkQuota(clientQuotaManager, "userA", "client1", 10, 1000, true);
+
+ // Remove <user> quota config, back to no quotas
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userA")),
+ Optional.empty(),
+ Optional.empty()
+ );
+ checkQuota(clientQuotaManager, "userA", "client1", Long.MAX_VALUE,
1000, false);
+ });
+ }
+
+ @Test
+ public void testSetAndRemoveUserClientQuota() {
+ // quotaTypesEnabled will be QuotaTypes.NoQuotas initially
+ withQuotaManager(QuotaType.PRODUCE, clientQuotaManager -> {
+ // Set <user, client-id> quota config
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userA")),
+ Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")),
+ Optional.of(new Quota(10, true))
+ );
+ checkQuota(clientQuotaManager, "userA", "client1", 10, 1000, true);
+
+ // Remove <user, client-id> quota config, back to no quotas
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userA")),
+ Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")),
+ Optional.empty()
+ );
+ checkQuota(clientQuotaManager, "userA", "client1", Long.MAX_VALUE,
1000, false);
+ });
+ }
+
+ @Test
+ public void testQuotaConfigPrecedence() {
+ withQuotaManager(QuotaType.PRODUCE, clientQuotaManager -> {
+ clientQuotaManager.updateQuota(
+ Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY),
+ Optional.empty(),
+ Optional.of(new Quota(1000, true))
+ );
+ clientQuotaManager.updateQuota(
+ Optional.empty(),
+ Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
+ Optional.of(new Quota(2000, true))
+ );
+ clientQuotaManager.updateQuota(
+ Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY),
+ Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
+ Optional.of(new Quota(3000, true))
+ );
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userA")),
+ Optional.empty(),
+ Optional.of(new Quota(4000, true))
+ );
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userA")),
+ Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")),
+ Optional.of(new Quota(5000, true))
+ );
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userB")),
+ Optional.empty(),
+ Optional.of(new Quota(6000, true))
+ );
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userB")),
+ Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")),
+ Optional.of(new Quota(7000, true))
+ );
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userB")),
+ Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
+ Optional.of(new Quota(8000, true))
+ );
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userC")),
+ Optional.empty(),
+ Optional.of(new Quota(10000, true))
+ );
+ clientQuotaManager.updateQuota(
+ Optional.empty(),
+ Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")),
+ Optional.of(new Quota(9000, true))
+ );
+
+ checkQuota(clientQuotaManager, "userA", "client1", 5000, 4500,
false); // <user, client> quota takes precedence over <user>
+ checkQuota(clientQuotaManager, "userA", "client2", 4000, 4500,
true); // <user> quota takes precedence over <client> and defaults
+ checkQuota(clientQuotaManager, "userA", "client3", 4000, 0, true);
// <user> quota is shared across clients of user
+ checkQuota(clientQuotaManager, "userA", "client1", 5000, 0,
false); // <user, client> is exclusive use, unaffected by other clients
+
+ checkQuota(clientQuotaManager, "userB", "client1", 7000, 8000,
true);
+ checkQuota(clientQuotaManager, "userB", "client2", 8000, 7000,
false); // Default per-client quota for exclusive use of <user, client>
+ checkQuota(clientQuotaManager, "userB", "client3", 8000, 7000,
false);
+
+ checkQuota(clientQuotaManager, "userD", "client1", 3000, 3500,
true); // Default <user, client> quota
+ checkQuota(clientQuotaManager, "userD", "client2", 3000, 2500,
false);
+ checkQuota(clientQuotaManager, "userE", "client1", 3000, 2500,
false);
+
+ // Remove default <user, client> quota config, revert to <user>
default
+ clientQuotaManager.updateQuota(
+ Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY),
+ Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
+ Optional.empty()
+ );
+ checkQuota(clientQuotaManager, "userD", "client1", 1000, 0,
false); // Metrics tags changed, restart counter
+ checkQuota(clientQuotaManager, "userE", "client4", 1000, 1500,
true);
+ checkQuota(clientQuotaManager, "userF", "client4", 1000, 800,
false); // Default <user> quota shared across clients of user
+ checkQuota(clientQuotaManager, "userF", "client5", 1000, 800,
true);
+
+ // Remove default <user> quota config, revert to <client-id>
default
+ clientQuotaManager.updateQuota(
+ Optional.of(ClientQuotaManager.DEFAULT_USER_ENTITY),
+ Optional.empty(),
+ Optional.empty()
+ );
+ checkQuota(clientQuotaManager, "userF", "client4", 2000, 0,
false); // Default <client-id> quota shared across client-id of all users
+ checkQuota(clientQuotaManager, "userF", "client5", 2000, 0, false);
+ checkQuota(clientQuotaManager, "userF", "client5", 2000, 2500,
true);
+ checkQuota(clientQuotaManager, "userG", "client5", 2000, 0, true);
+
+ // Update quotas
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userA")),
+ Optional.empty(),
+ Optional.of(new Quota(8000, true))
+ );
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userA")),
+ Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")),
+ Optional.of(new Quota(10000, true))
+ );
+ checkQuota(clientQuotaManager, "userA", "client2", 8000, 0, false);
+ checkQuota(clientQuotaManager, "userA", "client2", 8000, 4500,
true); // Throttled due to sum of new and earlier values
+ checkQuota(clientQuotaManager, "userA", "client1", 10000, 0,
false);
+ checkQuota(clientQuotaManager, "userA", "client1", 10000, 6000,
true);
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userA")),
+ Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")),
+ Optional.empty()
+ );
+ checkQuota(clientQuotaManager, "userA", "client6", 8000, 0, true);
// Throttled due to shared user quota
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userA")),
+ Optional.of(new
ClientQuotaManager.ClientIdEntity("client6")),
+ Optional.of(new Quota(11000, true))
+ );
+ checkQuota(clientQuotaManager, "userA", "client6", 11000, 8500,
false);
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userA")),
+ Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
+ Optional.of(new Quota(12000, true))
+ );
+ clientQuotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity("userA")),
+ Optional.of(new
ClientQuotaManager.ClientIdEntity("client6")),
+ Optional.empty()
+ );
+ checkQuota(clientQuotaManager, "userA", "client6", 12000, 4000,
true); // Throttled due to sum of new and earlier values
+
+ });
+ }
+
+ @Test
+ public void testQuotaViolation() {
+ withQuotaManager(QuotaType.PRODUCE, clientQuotaManager -> {
+ KafkaMetric queueSizeMetric =
metrics.metrics().get(metrics.metricName("queue-size", "Produce", ""));
+ clientQuotaManager.updateQuota(
+ Optional.empty(),
+ Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
+ Optional.of(new Quota(500, true))
+ );
+
+ // We have 10 seconds windows. Make sure that there is no quota
violation
+ // if we produce under the quota
+ for (int i = 0; i < 10; i++) {
+ assertEquals(0, maybeRecord(clientQuotaManager, "ANONYMOUS",
"unknown", 400));
+ time.sleep(1000);
+ }
+ assertEquals(0, ((Double)
queueSizeMetric.metricValue()).intValue());
+
+ // Create a spike.
+ // 400*10 + 2000 + 300 = 6300/10.5 = 600 bytes per second.
+ // (600 - quota)/quota*window-size = (600-500)/500*10.5 seconds =
2100
+ // 10.5 seconds because the last window is half complete
+ time.sleep(500);
+ int throttleTime = maybeRecord(clientQuotaManager, "ANONYMOUS",
"unknown", 2300);
+
+ assertEquals(2100, throttleTime, "Should be throttled");
+ throttle(clientQuotaManager, throttleTime, callback);
+ assertEquals(1, ((Double)
queueSizeMetric.metricValue()).intValue());
+ // After a request is delayed, the callback cannot be triggered
immediately
+ clientQuotaManager.processThrottledChannelReaperDoWork();
+ assertEquals(0, numCallbacks);
+ time.sleep(throttleTime);
+
+ // Callback can only be triggered after the delay time passes
+ clientQuotaManager.processThrottledChannelReaperDoWork();
+ assertEquals(0, ((Double)
queueSizeMetric.metricValue()).intValue());
+ assertEquals(1, numCallbacks);
+
+ // Could continue to see delays until the bursty sample disappears
+ for (int i = 0; i < 10; i++) {
+ maybeRecord(clientQuotaManager, "ANONYMOUS", "unknown", 400);
+ time.sleep(1000);
+ }
+
+ assertEquals(0, maybeRecord(clientQuotaManager, "ANONYMOUS",
"unknown", 0),
+ "Should be unthrottled since bursty sample has rolled
over");
+ });
+ }
+
+ @Test
+ public void testExpireThrottleTimeSensor() {
+ withQuotaManager(QuotaType.PRODUCE, clientQuotaManager -> {
+ clientQuotaManager.updateQuota(
+ Optional.empty(),
+ Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
+ Optional.of(new Quota(500, true))
+ );
+
+ maybeRecord(clientQuotaManager, "ANONYMOUS", "client1", 100);
+ // remove the throttle time sensor
+ metrics.removeSensor("ProduceThrottleTime-:client1");
+ // should not throw an exception even if the throttle time sensor
does not exist.
+ int throttleTime = maybeRecord(clientQuotaManager, "ANONYMOUS",
"client1", 10000);
+ assertTrue(throttleTime > 0, "Should be throttled");
+ // the sensor should get recreated
+ Sensor throttleTimeSensor =
metrics.getSensor("ProduceThrottleTime-:client1");
+ assertNotNull(throttleTimeSensor, "Throttle time sensor should
exist");
+ });
+ }
+
+ @Test
+ public void testExpireQuotaSensors() {
+ withQuotaManager(QuotaType.PRODUCE, clientQuotaManager -> {
+ clientQuotaManager.updateQuota(
+ Optional.empty(),
+ Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
+ Optional.of(new Quota(500, true))
+ );
+
+ maybeRecord(clientQuotaManager, "ANONYMOUS", "client1", 100);
+ // remove all the sensors
+ metrics.removeSensor("ProduceThrottleTime-:client1");
+ metrics.removeSensor("Produce-ANONYMOUS:client1");
+ // should not throw an exception
+ int throttleTime = maybeRecord(clientQuotaManager, "ANONYMOUS",
"client1", 10000);
+ assertTrue(throttleTime > 0, "Should be throttled");
+
+ // all the sensors should get recreated
+ Sensor throttleTimeSensor =
metrics.getSensor("ProduceThrottleTime-:client1");
+ assertNotNull(throttleTimeSensor, "Throttle time sensor should
exist");
+
+ Sensor byteRateSensor = metrics.getSensor("Produce-:client1");
+ assertNotNull(byteRateSensor, "Byte rate sensor should exist");
+ });
+ }
+
+ @Test
+ public void testClientIdNotSanitized() {
+ withQuotaManager(QuotaType.PRODUCE, clientQuotaManager -> {
+ String clientId = "client@#$%";
+ clientQuotaManager.updateQuota(
+ Optional.empty(),
+ Optional.of(ClientQuotaManager.DEFAULT_USER_CLIENT_ID),
+ Optional.of(new Quota(500, true))
+ );
+
+ maybeRecord(clientQuotaManager, "ANONYMOUS", clientId, 100);
+
+ // The metrics should use the raw client ID, even if the reporters
internally sanitize them
+ Sensor throttleTimeSensor =
metrics.getSensor("ProduceThrottleTime-:" + clientId);
+ assertNotNull(throttleTimeSensor, "Throttle time sensor should
exist");
+
+ Sensor byteRateSensor = metrics.getSensor("Produce-:" + clientId);
+ assertNotNull(byteRateSensor, "Byte rate sensor should exist");
+ });
+ }
+
+ @Test
+ public void testQuotaTypesEnabledUpdatesWithDefaultCallback() {
+ withQuotaManager(QuotaType.CONTROLLER_MUTATION, clientQuotaManager -> {
+ assertEquals(ClientQuotaManager.NO_QUOTAS,
clientQuotaManager.quotaTypesEnabled());
+ assertFalse(clientQuotaManager.quotasEnabled());
+
+ clientQuotaManager.updateQuota(Optional.empty(), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.of(new Quota(5, true)));
+ assertEquals(ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED,
clientQuotaManager.quotaTypesEnabled());
+ assertTrue(clientQuotaManager.quotasEnabled());
+
+ clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.empty(), Optional.of(new
Quota(5, true)));
+ assertEquals(ClientQuotaManager.USER_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED,
clientQuotaManager.quotaTypesEnabled());
+ assertTrue(clientQuotaManager.quotasEnabled());
+
+ clientQuotaManager.updateQuota(Optional.empty(), Optional.of(new
ClientQuotaManager.ClientIdEntity("client2")), Optional.of(new Quota(5, true)));
+ assertEquals(ClientQuotaManager.USER_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED,
clientQuotaManager.quotaTypesEnabled());
+ assertTrue(clientQuotaManager.quotasEnabled());
+
+ clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userB")), Optional.empty(), Optional.of(new
Quota(5, true)));
+ assertEquals(ClientQuotaManager.USER_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED,
clientQuotaManager.quotaTypesEnabled());
+ assertTrue(clientQuotaManager.quotasEnabled());
+
+ clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.of(new Quota(10,
true)));
+ assertEquals(ClientQuotaManager.USER_CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.USER_QUOTA_ENABLED, clientQuotaManager.quotaTypesEnabled());
+ assertTrue(clientQuotaManager.quotasEnabled());
+
+ clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.of(new Quota(12,
true)));
+ assertEquals(ClientQuotaManager.USER_CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.USER_QUOTA_ENABLED, clientQuotaManager.quotaTypesEnabled());
+ assertTrue(clientQuotaManager.quotasEnabled());
+
+ clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.empty(), Optional.empty());
+ assertEquals(ClientQuotaManager.USER_CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.USER_QUOTA_ENABLED, clientQuotaManager.quotaTypesEnabled());
+ assertTrue(clientQuotaManager.quotasEnabled());
+
+ clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userB")), Optional.empty(), Optional.empty());
+ assertEquals(ClientQuotaManager.USER_CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED,
clientQuotaManager.quotaTypesEnabled());
+ assertTrue(clientQuotaManager.quotasEnabled());
+
+ clientQuotaManager.updateQuota(Optional.empty(), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.empty());
+ assertEquals(ClientQuotaManager.USER_CLIENT_ID_QUOTA_ENABLED |
ClientQuotaManager.CLIENT_ID_QUOTA_ENABLED,
clientQuotaManager.quotaTypesEnabled());
+ assertTrue(clientQuotaManager.quotasEnabled());
+
+ clientQuotaManager.updateQuota(Optional.empty(), Optional.of(new
ClientQuotaManager.ClientIdEntity("client2")), Optional.empty());
+ assertEquals(ClientQuotaManager.USER_CLIENT_ID_QUOTA_ENABLED,
clientQuotaManager.quotaTypesEnabled());
+ assertTrue(clientQuotaManager.quotasEnabled());
+
+ clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.empty());
+ assertEquals(ClientQuotaManager.NO_QUOTAS,
clientQuotaManager.quotaTypesEnabled());
+ assertFalse(clientQuotaManager.quotasEnabled());
+
+ clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.empty());
+ assertEquals(ClientQuotaManager.NO_QUOTAS,
clientQuotaManager.quotaTypesEnabled());
+ assertFalse(clientQuotaManager.quotasEnabled());
+ });
+ }
+
+ @Test
+ public void testQuotaTypesEnabledUpdatesWithCustomCallback() {
+ ClientQuotaCallback customQuotaCallback = new ClientQuotaCallback() {
+ final Map<ClientQuotaEntity, Quota> quotas = new HashMap<>();
+ @Override
+ public void configure(Map<String, ?> configs) {}
+
+ @Override
+ public Map<String, String> quotaMetricTags(ClientQuotaType
quotaType, KafkaPrincipal principal, String clientId) {
+ return Map.of();
+ }
+
+ @Override
+ public Double quotaLimit(ClientQuotaType quotaType, Map<String,
String> metricTags) {
+ return 1.0;
+ }
+
+ @Override
+ public void updateQuota(ClientQuotaType quotaType,
ClientQuotaEntity entity, double newValue) {
+ quotas.put(entity, new Quota(newValue, true));
+ }
+
+ @Override
+ public void removeQuota(ClientQuotaType quotaType,
ClientQuotaEntity entity) {
+ quotas.remove(entity);
+ }
+
+ @Override
+ public boolean quotaResetRequired(ClientQuotaType quotaType) {
+ return false;
+ }
+
+ @Override
+ public void close() {}
+ };
+ ClientQuotaManager clientQuotaManager = new ClientQuotaManager(
+ new ClientQuotaManagerConfig(),
+ metrics,
+ QuotaType.CONTROLLER_MUTATION,
+ time,
+ "",
+ Optional.of(Plugin.wrapInstance(customQuotaCallback, metrics,
""))
+ );
+
+ try {
+ assertEquals(ClientQuotaManager.CUSTOM_QUOTAS,
clientQuotaManager.quotaTypesEnabled());
+ assertTrue(clientQuotaManager.quotasEnabled(), "quotasEnabled
should be true with custom callback");
+
+ clientQuotaManager.updateQuota(Optional.empty(), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.of(new Quota(12,
true)));
+ assertEquals(ClientQuotaManager.CUSTOM_QUOTAS,
clientQuotaManager.quotaTypesEnabled());
+
+ clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.empty(), Optional.of(new
Quota(12, true)));
+ assertEquals(ClientQuotaManager.CUSTOM_QUOTAS,
clientQuotaManager.quotaTypesEnabled());
+ assertTrue(clientQuotaManager.quotasEnabled(), "quotasEnabled
should remain true");
+
+ clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.of(new Quota(12,
true)));
+ assertEquals(ClientQuotaManager.CUSTOM_QUOTAS,
clientQuotaManager.quotaTypesEnabled());
+ assertTrue(clientQuotaManager.quotasEnabled(), "quotasEnabled
should remain true");
+
+ clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.empty());
+ clientQuotaManager.updateQuota(Optional.of(new
ClientQuotaManager.UserEntity("userA")), Optional.empty(), Optional.empty());
+ clientQuotaManager.updateQuota(Optional.empty(), Optional.of(new
ClientQuotaManager.ClientIdEntity("client1")), Optional.empty());
+ assertEquals(ClientQuotaManager.CUSTOM_QUOTAS,
clientQuotaManager.quotaTypesEnabled());
+ assertTrue(clientQuotaManager.quotasEnabled(), "quotasEnabled
should remain true");
+ } finally {
+ clientQuotaManager.shutdown();
+ }
+ }
+
+ record UserClient(
+ String user,
+ String clientId,
+ Optional<ClientQuotaEntity.ConfigEntity> configUser,
+ Optional<ClientQuotaEntity.ConfigEntity> configClientEntity) { }
+
+}
diff --git
a/server/src/test/java/org/apache/kafka/server/quota/ClientRequestQuotaManagerTest.java
b/server/src/test/java/org/apache/kafka/server/quota/ClientRequestQuotaManagerTest.java
new file mode 100644
index 00000000000..8ce603b6b27
--- /dev/null
+++
b/server/src/test/java/org/apache/kafka/server/quota/ClientRequestQuotaManagerTest.java
@@ -0,0 +1,104 @@
+/*
+ * 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 org.apache.kafka.server.quota;
+
+import org.apache.kafka.common.metrics.KafkaMetric;
+import org.apache.kafka.common.metrics.Quota;
+import org.apache.kafka.server.config.ClientQuotaManagerConfig;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Optional;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class ClientRequestQuotaManagerTest extends BaseClientQuotaManagerTest {
+ private final ClientQuotaManagerConfig config = new
ClientQuotaManagerConfig();
+
+ private static double millisToPercent(double millis) {
+ return millis * 1000 * 1000 *
ClientRequestQuotaManager.NANOS_TO_PERCENTAGE_PER_SECOND;
+ }
+
+ @Test
+ public void testRequestPercentageQuotaViolation() {
+ ClientRequestQuotaManager clientRequestQuotaManager = new
ClientRequestQuotaManager(config, metrics, time, "", Optional.empty());
+ ClientQuotaEntity.ConfigEntity userEntity = new
ClientQuotaManager.UserEntity("ANONYMOUS");
+ ClientQuotaEntity.ConfigEntity clientEntity = new
ClientQuotaManager.ClientIdEntity("test-client");
+
+ clientRequestQuotaManager.updateQuota(
+ Optional.of(userEntity),
+ Optional.of(clientEntity),
+ Optional.of(Quota.upperBound(1))
+ );
+ KafkaMetric queueSizeMetric =
metrics.metrics().get(metrics.metricName("queue-size",
QuotaType.REQUEST.toString(), ""));
+
+ try {
+ // We have 10 seconds windows. Make sure that there is no quota
violation
+ // if we are under the quota
+ for (int i = 0; i < 10; i++) {
+ assertEquals(0, maybeRecord(clientRequestQuotaManager,
"ANONYMOUS", "test-client", millisToPercent(4)));
+ time.sleep(1000);
+ }
+ assertEquals(0, ((Double)
queueSizeMetric.metricValue()).intValue());
+
+ // Create a spike.
+ // quota = 1% (10ms per second)
+ // 4*10 + 67.1 = 107.1/10.5 = 10.2ms per second.
+ // (10.2 - quota)/quota*window-size = (10.2-10)/10*10.5 seconds =
210ms
+ // 10.5 seconds interval because the last window is half complete
+ time.sleep(500);
+ int throttleTime = maybeRecord(clientRequestQuotaManager,
"ANONYMOUS", "test-client", millisToPercent(67.1));
+
+ assertEquals(210, throttleTime, "Should be throttled");
+
+ throttle(clientRequestQuotaManager, throttleTime, callback);
+ assertEquals(1, ((Double)
queueSizeMetric.metricValue()).intValue());
+ // After a request is delayed, the callback cannot be triggered
immediately
+ clientRequestQuotaManager.processThrottledChannelReaperDoWork();
+ assertEquals(0, numCallbacks);
+ time.sleep(throttleTime);
+
+ // Callback can only be triggered after the delay time passes
+ clientRequestQuotaManager.processThrottledChannelReaperDoWork();
+ assertEquals(0, ((Double)
queueSizeMetric.metricValue()).intValue());
+ assertEquals(1, numCallbacks);
+
+ // Could continue to see delays until the bursty sample disappears
+ for (int i = 0; i < 11; i++) {
+ maybeRecord(clientRequestQuotaManager, "ANONYMOUS",
"test-client", millisToPercent(4));
+ time.sleep(1000);
+ }
+
+ assertEquals(0,
+ maybeRecord(clientRequestQuotaManager, "ANONYMOUS",
"test-client", 0), "Should be unthrottled since bursty sample has rolled over");
+
+ // Create a very large spike which requires > one quota window to
bring within quota
+ assertEquals(1000, maybeRecord(clientRequestQuotaManager,
"ANONYMOUS", "test-client", millisToPercent(500)));
+ for (int i = 0; i < 10; i++) {
+ time.sleep(1000);
+ assertEquals(1000, maybeRecord(clientRequestQuotaManager,
"ANONYMOUS", "test-client", 0));
+ }
+ time.sleep(1000);
+ assertEquals(0,
+ maybeRecord(clientRequestQuotaManager, "ANONYMOUS",
"test-client", 0), "Should be unthrottled since bursty sample has rolled over");
+
+ } finally {
+ clientRequestQuotaManager.shutdown();
+ }
+ }
+
+}
diff --git
a/server/src/test/java/org/apache/kafka/server/quota/ControllerMutationQuotaManagerTest.java
b/server/src/test/java/org/apache/kafka/server/quota/ControllerMutationQuotaManagerTest.java
new file mode 100644
index 00000000000..b9e2c1219bb
--- /dev/null
+++
b/server/src/test/java/org/apache/kafka/server/quota/ControllerMutationQuotaManagerTest.java
@@ -0,0 +1,243 @@
+/*
+ * 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 org.apache.kafka.server.quota;
+
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.errors.ThrottlingQuotaExceededException;
+import org.apache.kafka.common.metrics.KafkaMetric;
+import org.apache.kafka.common.metrics.MetricConfig;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.Quota;
+import org.apache.kafka.common.metrics.QuotaViolationException;
+import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.metrics.stats.TokenBucket;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.server.config.ClientQuotaManagerConfig;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
+import static
org.apache.kafka.server.quota.ControllerMutationQuotaManager.throttleTimeMs;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class ControllerMutationQuotaManagerTest extends
BaseClientQuotaManagerTest {
+
+ private static final String USER = "ANONYMOUS";
+ private static final String CLIENT_ID = "test-client";
+
+ private final ClientQuotaManagerConfig config = new
ClientQuotaManagerConfig(10, 1);
+
+ @Test
+ public void testStrictControllerMutationQuotaViolation() {
+ MockTime time = new MockTime(0, System.currentTimeMillis(), 0);
+ try (Metrics metrics = new Metrics(time)) {
+ Sensor sensor = metrics.sensor("sensor", new MetricConfig()
+ .quota(Quota.upperBound(10))
+ .timeWindow(1, TimeUnit.SECONDS)
+ .samples(10));
+ MetricName metricName = metrics.metricName("rate", "test-group");
+ assertTrue(sensor.add(metricName, new TokenBucket()));
+
+ StrictControllerMutationQuota quota = new
StrictControllerMutationQuota(time, sensor);
+ assertFalse(quota.isExceeded());
+
+ // Recording a first value at T to bring the tokens to 10. Value
is accepted
+ // because the quota is not exhausted yet.
+ quota.record(90);
+ assertFalse(quota.isExceeded());
+ assertEquals(0, quota.throttleTime());
+
+ // Recording a second value at T to bring the tokens to -80. Value
is accepted
+ quota.record(90);
+ assertFalse(quota.isExceeded());
+ assertEquals(0, quota.throttleTime());
+
+ // Recording a third value at T is rejected immediately because
there are no
+ // tokens available in the bucket.
+ assertThrows(ThrottlingQuotaExceededException.class, () ->
quota.record(90));
+ assertTrue(quota.isExceeded());
+ assertEquals(8000, quota.throttleTime());
+
+ // Throttle time is adjusted with time
+ time.sleep(5000);
+ assertEquals(3000, quota.throttleTime());
+ }
+ }
+
+ @Test
+ public void testPermissiveControllerMutationQuotaViolation() {
+ MockTime time = new MockTime(0, System.currentTimeMillis(), 0);
+ try (Metrics metrics = new Metrics(time)) {
+ Sensor sensor = metrics.sensor("sensor", new MetricConfig()
+ .quota(Quota.upperBound(10))
+ .timeWindow(1, TimeUnit.SECONDS)
+ .samples(10));
+ MetricName metricName = metrics.metricName("rate", "test-group");
+ assertTrue(sensor.add(metricName, new TokenBucket()));
+
+ PermissiveControllerMutationQuota quota = new
PermissiveControllerMutationQuota(time, sensor);
+ assertFalse(quota.isExceeded());
+
+ // Recording a first value at T to bring the tokens 10. Value is
accepted
+ // because the quota is not exhausted yet.
+ quota.record(90);
+ assertFalse(quota.isExceeded());
+ assertEquals(0, quota.throttleTime());
+
+ // Recording a second value at T to bring the tokens to -80. Value
is accepted
+ quota.record(90);
+ assertFalse(quota.isExceeded());
+ assertEquals(8000, quota.throttleTime());
+
+ // Recording a third value at T to bring the tokens to -170. Value
is accepted
+ // even though the quota is exhausted.
+ quota.record(90);
+ assertFalse(quota.isExceeded()); // quota is never exceeded
+ assertEquals(17000, quota.throttleTime());
+
+ // Throttle time is adjusted with time
+ time.sleep(5000);
+ assertEquals(12000, quota.throttleTime());
+ }
+ }
+
+ private void withQuotaManager(Consumer<ControllerMutationQuotaManager> f) {
+ ControllerMutationQuotaManager quotaManager = new
ControllerMutationQuotaManager(config, metrics, time, "", Optional.empty());
+ try {
+ f.accept(quotaManager);
+ } finally {
+ quotaManager.shutdown();
+ }
+ }
+
+ @Test
+ public void testThrottleTime() {
+ MockTime time = new MockTime(0, System.currentTimeMillis(), 0);
+ try (Metrics metrics = new Metrics(time)) {
+ Sensor sensor = metrics.sensor("sensor");
+ MetricName metricName = metrics.metricName("tokens", "test-group");
+ sensor.add(metricName, new TokenBucket());
+ KafkaMetric metric = metrics.metric(metricName);
+
+ assertEquals(0, throttleTimeMs(new QuotaViolationException(metric,
0, 10)));
+ assertEquals(500, throttleTimeMs(new
QuotaViolationException(metric, -5, 10)));
+ assertEquals(1000, throttleTimeMs(new
QuotaViolationException(metric, -10, 10)));
+ }
+ }
+
+ @Test
+ public void testControllerMutationQuotaViolation() {
+ withQuotaManager(quotaManager -> {
+ quotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity(USER)),
+ Optional.of(new
ClientQuotaManager.ClientIdEntity(CLIENT_ID)),
+ Optional.of(Quota.upperBound(10))
+ );
+ KafkaMetric queueSizeMetric = metrics.metrics().get(
+ metrics.metricName("queue-size",
QuotaType.CONTROLLER_MUTATION.toString(), ""));
+
+ // Verify that there is no quota violation if we remain under the
quota.
+ for (int i = 0; i < 10; i++) {
+ assertEquals(0, maybeRecord(quotaManager, USER, CLIENT_ID,
10));
+ time.sleep(1000);
+ }
+ assertEquals(0, ((Double)
queueSizeMetric.metricValue()).intValue());
+
+ // Create a spike worth of 110 mutations.
+ // Current tokens in the bucket = 100
+ // As we use the Strict enforcement, the quota is checked before
updating the rate. Hence,
+ // the spike is accepted and no quota violation error is raised.
+ int throttleTime = maybeRecord(quotaManager, USER, CLIENT_ID, 110);
+ assertEquals(0, throttleTime, "Should not be throttled");
+
+ // Create a spike worth of 110 mutations.
+ // Current tokens in the bucket = 100 - 110 = -10
+ // As the quota is already violated, the spike is rejected
immediately without updating the
+ // rate. The client must wait:
+ // 10 / 10 = 1s
+ throttleTime = maybeRecord(quotaManager, USER, CLIENT_ID, 110);
+ assertEquals(1000, throttleTime, "Should be throttled");
+
+ // Throttle
+ throttle(quotaManager, throttleTime, callback);
+ assertEquals(1, ((Double)
queueSizeMetric.metricValue()).intValue());
+
+ // After a request is delayed, the callback cannot be triggered
immediately
+ quotaManager.processThrottledChannelReaperDoWork();
+ assertEquals(0, numCallbacks);
+
+ // Callback can only be triggered after the delay time passes
+ time.sleep(throttleTime);
+ quotaManager.processThrottledChannelReaperDoWork();
+ assertEquals(0, ((Double)
queueSizeMetric.metricValue()).intValue());
+ assertEquals(1, numCallbacks);
+
+ // Retry to spike worth of 110 mutations after having waited the
required throttle time.
+ // Current tokens in the bucket = 0
+ throttleTime = maybeRecord(quotaManager, USER, CLIENT_ID, 110);
+ assertEquals(0, throttleTime, "Should not be throttled");
+ });
+ }
+
+ @Test
+ public void
testNewStrictQuotaForReturnsUnboundedQuotaWhenQuotaIsDisabled() {
+ withQuotaManager(quotaManager ->
+
assertEquals(ControllerMutationQuota.UNBOUNDED_CONTROLLER_MUTATION_QUOTA,
+ quotaManager.newStrictQuotaFor(buildSession(USER),
CLIENT_ID)));
+ }
+
+ @Test
+ public void testNewStrictQuotaForReturnsStrictQuotaWhenQuotaIsEnabled() {
+ withQuotaManager(quotaManager -> {
+ quotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity(USER)),
+ Optional.of(new
ClientQuotaManager.ClientIdEntity(CLIENT_ID)),
+ Optional.of(Quota.upperBound(10))
+ );
+ ControllerMutationQuota quota =
quotaManager.newStrictQuotaFor(buildSession(USER), CLIENT_ID);
+ assertInstanceOf(StrictControllerMutationQuota.class, quota);
+ });
+ }
+
+ @Test
+ public void
testNewPermissiveQuotaForReturnsUnboundedQuotaWhenQuotaIsDisabled() {
+ withQuotaManager(quotaManager ->
+
assertEquals(ControllerMutationQuota.UNBOUNDED_CONTROLLER_MUTATION_QUOTA,
+ quotaManager.newPermissiveQuotaFor(buildSession(USER),
CLIENT_ID)));
+ }
+
+ @Test
+ public void
testNewPermissiveQuotaForReturnsStrictQuotaWhenQuotaIsEnabled() {
+ withQuotaManager(quotaManager -> {
+ quotaManager.updateQuota(
+ Optional.of(new ClientQuotaManager.UserEntity(USER)),
+ Optional.of(new
ClientQuotaManager.ClientIdEntity(CLIENT_ID)),
+ Optional.of(Quota.upperBound(10))
+ );
+ ControllerMutationQuota quota =
quotaManager.newPermissiveQuotaFor(buildSession(USER), CLIENT_ID);
+ assertInstanceOf(PermissiveControllerMutationQuota.class, quota);
+ });
+ }
+
+}
