| /** |
| * 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.{ConcurrentHashMap, DelayQueue, TimeUnit} |
| |
| import kafka.utils.{ShutdownableThread, Logging} |
| import org.apache.kafka.common.MetricName |
| import org.apache.kafka.common.metrics._ |
| import org.apache.kafka.common.metrics.stats.{Total, Rate, Avg} |
| import java.util.concurrent.locks.ReentrantReadWriteLock |
| |
| import org.apache.kafka.common.utils.Time |
| |
| /** |
| * Represents the sensors aggregated per client |
| * @param quotaSensor @Sensor that tracks the quota |
| * @param throttleTimeSensor @Sensor that tracks the throttle time |
| */ |
| private case class ClientSensors(quotaSensor: Sensor, throttleTimeSensor: Sensor) |
| |
| /** |
| * Configuration settings for quota management |
| * @param quotaBytesPerSecondDefault The default bytes per second quota allocated to any client |
| * @param numQuotaSamples The number of samples to retain in memory |
| * @param quotaWindowSizeSeconds The time span of each sample |
| * |
| */ |
| case class ClientQuotaManagerConfig(quotaBytesPerSecondDefault: Long = |
| ClientQuotaManagerConfig.QuotaBytesPerSecondDefault, |
| numQuotaSamples: Int = |
| ClientQuotaManagerConfig.DefaultNumQuotaSamples, |
| quotaWindowSizeSeconds: Int = |
| ClientQuotaManagerConfig.DefaultQuotaWindowSizeSeconds) |
| |
| object ClientQuotaManagerConfig { |
| val QuotaBytesPerSecondDefault = Long.MaxValue |
| // Always have 10 whole windows + 1 current window |
| val DefaultNumQuotaSamples = 11 |
| val DefaultQuotaWindowSizeSeconds = 1 |
| // Purge sensors after 1 hour of inactivity |
| val InactiveSensorExpirationTimeSeconds = 3600 |
| } |
| |
| /** |
| * Helper class that records per-client metrics. It is also responsible for maintaining Quota usage statistics |
| * for all clients. |
| * @param config @ClientQuotaManagerConfig quota configs |
| * @param metrics @Metrics Metrics instance |
| * @param apiKey API Key for the request |
| * @param time @Time object to use |
| */ |
| class ClientQuotaManager(private val config: ClientQuotaManagerConfig, |
| private val metrics: Metrics, |
| private val apiKey: String, |
| private val time: Time) extends Logging { |
| private val overriddenQuota = new ConcurrentHashMap[String, Quota]() |
| private val defaultQuota = Quota.upperBound(config.quotaBytesPerSecondDefault) |
| private val lock = new ReentrantReadWriteLock() |
| private val delayQueue = new DelayQueue[ThrottledResponse]() |
| val throttledRequestReaper = new ThrottledRequestReaper(delayQueue) |
| throttledRequestReaper.start() |
| |
| private val delayQueueSensor = metrics.sensor(apiKey + "-delayQueue") |
| delayQueueSensor.add(metrics.metricName("queue-size", |
| apiKey, |
| "Tracks the size of the delay queue"), new Total()) |
| |
| /** |
| * Reaper thread that triggers callbacks on all throttled requests |
| * @param delayQueue DelayQueue to dequeue from |
| */ |
| class ThrottledRequestReaper(delayQueue: DelayQueue[ThrottledResponse]) extends ShutdownableThread( |
| "ThrottledRequestReaper-%s".format(apiKey), false) { |
| |
| override def doWork(): Unit = { |
| val response: ThrottledResponse = delayQueue.poll(1, TimeUnit.SECONDS) |
| if (response != null) { |
| // Decrement the size of the delay queue |
| delayQueueSensor.record(-1) |
| trace("Response throttled for: " + response.throttleTimeMs + " ms") |
| response.execute() |
| } |
| } |
| } |
| |
| /** |
| * Records that a clientId changed some metric being throttled (produced/consumed bytes, QPS etc.) |
| * @param clientId clientId that produced the data |
| * @param value amount of data written in bytes |
| * @param callback Callback function. This will be triggered immediately if quota is not violated. |
| * If there is a quota violation, this callback will be triggered after a delay |
| * @return Number of milliseconds to delay the response in case of Quota violation. |
| * Zero otherwise |
| */ |
| def recordAndMaybeThrottle(clientId: String, value: Int, callback: Int => Unit): Int = { |
| val clientSensors = getOrCreateQuotaSensors(clientId) |
| var throttleTimeMs = 0 |
| try { |
| clientSensors.quotaSensor.record(value) |
| // trigger the callback immediately if quota is not violated |
| callback(0) |
| } catch { |
| case qve: QuotaViolationException => |
| // Compute the delay |
| val clientMetric = metrics.metrics().get(clientRateMetricName(clientId)) |
| throttleTimeMs = throttleTime(clientMetric, getQuotaMetricConfig(quota(clientId))) |
| clientSensors.throttleTimeSensor.record(throttleTimeMs) |
| // If delayed, add the element to the delayQueue |
| delayQueue.add(new ThrottledResponse(time, throttleTimeMs, callback)) |
| delayQueueSensor.record() |
| logger.debug("Quota violated for sensor (%s). Delay time: (%d)".format(clientSensors.quotaSensor.name(), throttleTimeMs)) |
| } |
| throttleTimeMs |
| } |
| |
| /* |
| * This calculates the amount of time needed to bring the metric within quota |
| * assuming that no new metrics are recorded. |
| * |
| * Basically, if O is the observed rate and T is the target rate over a window of W, to bring O down to T, |
| * we need to add a delay of X to W such that O * W / (W + X) = T. |
| * Solving for X, we get X = (O - T)/T * W. |
| */ |
| private def throttleTime(clientMetric: KafkaMetric, config: MetricConfig): Int = { |
| val rateMetric: Rate = measurableAsRate(clientMetric.metricName(), clientMetric.measurable()) |
| val quota = config.quota() |
| val difference = clientMetric.value() - quota.bound |
| // Use the precise window used by the rate calculation |
| val throttleTimeMs = difference / quota.bound * rateMetric.windowSize(config, time.milliseconds()) |
| throttleTimeMs.round.toInt |
| } |
| |
| // Casting to Rate because we only use Rate in Quota computation |
| private def measurableAsRate(name: MetricName, measurable: Measurable): Rate = { |
| measurable match { |
| case r: Rate => r |
| case _ => throw new IllegalArgumentException(s"Metric $name is not a Rate metric, value $measurable") |
| } |
| } |
| |
| /** |
| * Returns the quota for the specified clientId |
| */ |
| def quota(clientId: String): Quota = |
| if (overriddenQuota.containsKey(clientId)) overriddenQuota.get(clientId) else defaultQuota |
| |
| /* |
| * This function either returns the sensors for a given client id or creates them if they don't exist |
| * First sensor of the tuple is the quota enforcement sensor. Second one is the throttle time sensor |
| */ |
| private def getOrCreateQuotaSensors(clientId: String): ClientSensors = { |
| |
| // Names of the sensors to access |
| val quotaSensorName = getQuotaSensorName(clientId) |
| val throttleTimeSensorName = getThrottleTimeSensorName(clientId) |
| var quotaSensor: Sensor = null |
| var throttleTimeSensor: Sensor = null |
| |
| /* Acquire the read lock to fetch the sensors. It is safe to call getSensor from multiple threads. |
| * The read lock allows a thread to create a sensor in isolation. The thread creating the sensor |
| * will acquire the write lock and prevent the sensors from being read while they are being created. |
| * It should be sufficient to simply check if the sensor is null without acquiring a read lock but the |
| * sensor being present doesn't mean that it is fully initialized i.e. all the Metrics may not have been added. |
| * This read lock waits until the writer thread has released its lock i.e. fully initialized the sensor |
| * at which point it is safe to read |
| */ |
| lock.readLock().lock() |
| try { |
| quotaSensor = metrics.getSensor(quotaSensorName) |
| throttleTimeSensor = metrics.getSensor(throttleTimeSensorName) |
| } |
| finally { |
| lock.readLock().unlock() |
| } |
| |
| /* If the sensor is null, try to create it else return the created sensor |
| * Either of the sensors can be null, hence null checks on both |
| */ |
| if (quotaSensor == null || throttleTimeSensor == null) { |
| /* Acquire a write lock because the sensor may not have been created and we only want one thread to create it. |
| * Note that multiple threads may acquire the write lock if they all see a null sensor initially |
| * In this case, the writer checks the sensor after acquiring the lock again. |
| * This is safe from Double Checked Locking because the references are read |
| * after acquiring read locks and hence they cannot see a partially published reference |
| */ |
| lock.writeLock().lock() |
| try { |
| // Set the var for both sensors in case another thread has won the race to acquire the write lock. This will |
| // ensure that we initialise `ClientSensors` with non-null parameters. |
| quotaSensor = metrics.getSensor(quotaSensorName) |
| throttleTimeSensor = metrics.getSensor(throttleTimeSensorName) |
| if (throttleTimeSensor == null) { |
| // create the throttle time sensor also. Use default metric config |
| throttleTimeSensor = metrics.sensor(throttleTimeSensorName, |
| null, |
| ClientQuotaManagerConfig.InactiveSensorExpirationTimeSeconds) |
| throttleTimeSensor.add(metrics.metricName("throttle-time", |
| apiKey, |
| "Tracking average throttle-time per client", |
| "client-id", |
| clientId), new Avg()) |
| } |
| |
| |
| if (quotaSensor == null) { |
| quotaSensor = metrics.sensor(quotaSensorName, |
| getQuotaMetricConfig(quota(clientId)), |
| ClientQuotaManagerConfig.InactiveSensorExpirationTimeSeconds) |
| quotaSensor.add(clientRateMetricName(clientId), new Rate()) |
| } |
| } finally { |
| lock.writeLock().unlock() |
| } |
| } |
| // return the read or created sensors |
| ClientSensors(quotaSensor, throttleTimeSensor) |
| } |
| |
| private def getThrottleTimeSensorName(clientId: String): String = apiKey + "ThrottleTime-" + clientId |
| |
| private def getQuotaSensorName(clientId: String): String = apiKey + "-" + clientId |
| |
| private def getQuotaMetricConfig(quota: Quota): MetricConfig = { |
| new MetricConfig() |
| .timeWindow(config.quotaWindowSizeSeconds, TimeUnit.SECONDS) |
| .samples(config.numQuotaSamples) |
| .quota(quota) |
| } |
| |
| /** |
| * Reset quotas to the default value for the given clientId |
| * @param clientId client to override |
| */ |
| def resetQuota(clientId: String) = { |
| updateQuota(clientId, defaultQuota) |
| } |
| |
| /** |
| * Overrides quotas per clientId |
| * @param clientId client to override |
| * @param quota custom quota to apply |
| */ |
| def updateQuota(clientId: String, quota: Quota) = { |
| /* |
| * Acquire the write lock to apply changes in the quota objects. |
| * This method changes the quota in the overriddenQuota map and applies the update on the actual KafkaMetric object (if it exists). |
| * If the KafkaMetric hasn't been created, the most recent value will be used from the overriddenQuota map. |
| * The write lock prevents quota update and creation at the same time. It also guards against concurrent quota change |
| * notifications |
| */ |
| lock.writeLock().lock() |
| try { |
| logger.info(s"Changing quota for clientId $clientId to ${quota.bound()}") |
| |
| if (quota.equals(defaultQuota)) |
| this.overriddenQuota.remove(clientId) |
| else |
| this.overriddenQuota.put(clientId, quota) |
| |
| // Change the underlying metric config if the sensor has been created |
| val allMetrics = metrics.metrics() |
| val quotaMetricName = clientRateMetricName(clientId) |
| if (allMetrics.containsKey(quotaMetricName)) { |
| logger.info(s"Sensor for clientId $clientId already exists. Changing quota to ${quota.bound()} in MetricConfig") |
| allMetrics.get(quotaMetricName).config(getQuotaMetricConfig(quota)) |
| } |
| } finally { |
| lock.writeLock().unlock() |
| } |
| } |
| |
| private def clientRateMetricName(clientId: String): MetricName = { |
| metrics.metricName("byte-rate", apiKey, |
| "Tracking byte-rate per client", |
| "client-id", clientId) |
| } |
| |
| def shutdown() = { |
| throttledRequestReaper.shutdown() |
| } |
| } |
