Repository: apache/kafka
File: core/src/main/scala/kafka/server/ClientQuotaManager.scala (Scala)

Code before:
 * Quota limits including defaults may be updated dynamically. The implementation is optimized for the case
 * where a single level of quotas is configured.
 *
 * @param config @ClientQuotaManagerConfig quota configs
 * @param metrics @Metrics Metrics instance
 * @param quotaType Quota type of this quota manager
 * @param time @Time object to use
 * @param threadNamePrefix The thread prefix to use
 * @param clientQuotaCallbackPlugin An optional @ClientQuotaCallback and 
 *                                  wrap it in a {@link org.apache.kafka.common.internals.Plugin}
 */
class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
                         private val metrics: Metrics,
                         private val quotaType: QuotaType,
                         private val time: Time,
                         private val threadNamePrefix: String,
                         private val clientQuotaCallbackPlugin: Option[Plugin[ClientQuotaCallback]] = None) extends Logging {

  private val lock = new ReentrantReadWriteLock()
  private val sensorAccessor = new SensorAccess(lock, metrics)
  private val quotaCallback = clientQuotaCallbackPlugin match {
    case Some(plugin) => plugin.get()
    case None => new DefaultQuotaCallback
  }
  private val clientQuotaType = QuotaType.toClientQuotaType(quotaType)
  private val activeQuotaEntities = new util.HashSet[KafkaQuotaEntity]()

  @volatile
  var quotaTypesEnabled = clientQuotaCallbackPlugin match {
    case Some(_) => QuotaTypes.CustomQuotas
    case None => QuotaTypes.NoQuotas
  }

  private val delayQueueSensor = metrics.sensor(quotaType.toString + "-delayQueue")
  delayQueueSensor.add(metrics.metricName("queue-size", quotaType.toString,
    "Tracks the size of the delay queue"), new CumulativeSum())

  private val delayQueue = new DelayQueue[ThrottledChannel]()
  private[server] val throttledChannelReaper = new ThrottledChannelReaper(delayQueue, threadNamePrefix)
  start() // Use start method to keep spotbugs happy
  private def start(): Unit = {
    throttledChannelReaper.start()
  }

  /**
   * Reaper thread that triggers channel unmute callbacks on all throttled channels
   * @param delayQueue DelayQueue to dequeue from
   */
  class ThrottledChannelReaper(delayQueue: DelayQueue[ThrottledChannel], prefix: String) extends ShutdownableThread(
    s"${prefix}ThrottledChannelReaper-$quotaType", false) {


Review comment (performance):
Do we really need to maintain this new set of quota entities? When custom quotas are used, I am not sure we can optimize based on active entities anyway. And for the default quota callback, we already have the map containing all entities, couldn't we use that instead?

Code after:
 * Quota limits including defaults may be updated dynamically. The implementation is optimized for the case
 * where a single level of quotas is configured.
 *
 * @param config @ClientQuotaManagerConfig quota configs
 * @param metrics @Metrics Metrics instance
 * @param quotaType Quota type of this quota manager
 * @param time @Time object to use
 * @param threadNamePrefix The thread prefix to use
 * @param clientQuotaCallbackPlugin An optional @ClientQuotaCallback and 
 *                                  wrap it in a {@link org.apache.kafka.common.internals.Plugin}
 */
class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
                         private val metrics: Metrics,
                         private val quotaType: QuotaType,
                         private val time: Time,
                         private val threadNamePrefix: String,
                         private val clientQuotaCallbackPlugin: Option[Plugin[ClientQuotaCallback]] = None) extends Logging {

  private val lock = new ReentrantReadWriteLock()
  private val sensorAccessor = new SensorAccess(lock, metrics)
  private val quotaCallback = clientQuotaCallbackPlugin match {
    case Some(plugin) => plugin.get()
    case None => new DefaultQuotaCallback
  }
  private val clientQuotaType = QuotaType.toClientQuotaType(quotaType)
  private val activeQuotaEntities = new ConcurrentHashMap[Int, Int]() // Key is QuotaTypes, value is the number of active quota entities of that type



  @volatile
  var quotaTypesEnabled = clientQuotaCallbackPlugin match {
    case Some(_) => QuotaTypes.CustomQuotas
    case None => QuotaTypes.NoQuotas
  }

  private val delayQueueSensor = metrics.sensor(quotaType.toString + "-delayQueue")
  delayQueueSensor.add(metrics.metricName("queue-size", quotaType.toString,
    "Tracks the size of the delay queue"), new CumulativeSum())

  private val delayQueue = new DelayQueue[ThrottledChannel]()
  private[server] val throttledChannelReaper = new ThrottledChannelReaper(delayQueue, threadNamePrefix)
  start() // Use start method to keep spotbugs happy
  private def start(): Unit = {
    throttledChannelReaper.start()
  }

  /**
   * Reaper thread that triggers channel unmute callbacks on all throttled channels
   * @param delayQueue DelayQueue to dequeue from
   */
  class ThrottledChannelReaper(delayQueue: DelayQueue[ThrottledChannel], prefix: String) extends ShutdownableThread(