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

Code before:
    }
    val previousQuotaTypesEnabled = quotaTypesEnabled

    // Update activeQuotaTypes counts
    activeQuotaTypes.compute(quotaTypeKey, (_, count) =>
      if (increment) Option(count).getOrElse(0) + 1
      else if (Option(count).exists(_ > 1)) count - 1
      else 0
    )

    // Update quotaTypesEnabled based on activeQuotaTypes counts
    quotaTypesEnabled = clientQuotaCallbackPlugin match {
      case Some(_) => QuotaTypes.CustomQuotas
      case None =>
        var newQuotaTypes = QuotaTypes.NoQuotas
        if (activeQuotaTypes.getOrDefault(QuotaTypes.UserClientIdQuotaEnabled, 0) > 0)
          newQuotaTypes |= QuotaTypes.UserClientIdQuotaEnabled
        if (activeQuotaTypes.getOrDefault(QuotaTypes.UserQuotaEnabled, 0) > 0)
          newQuotaTypes |= QuotaTypes.UserQuotaEnabled
        if (activeQuotaTypes.getOrDefault(QuotaTypes.ClientIdQuotaEnabled, 0) > 0)
          newQuotaTypes |= QuotaTypes.ClientIdQuotaEnabled
        newQuotaTypes
    }

    if (previousQuotaTypesEnabled != quotaTypesEnabled) {
      debug(s"Quota types enabled changed from $previousQuotaTypesEnabled to $quotaTypesEnabled")
    }
  }

  /**
   * Updates metrics configs. This is invoked when quota configs are updated when partitions leaders change 
   * and custom callbacks that implement partition-based quotas have updated quotas.
   *
   * @param updatedQuotaEntity If set to one entity and quotas have only been enabled at one
   *    level, then an optimized update is performed with a single metric update. If None is provided,
   *    or if custom callbacks are used or if multi-level quotas have been enabled, all metric configs
   *    are checked and updated if required.
   */
  def updateQuotaMetricConfigs(updatedQuotaEntity: Option[KafkaQuotaEntity] = None): Unit = {
    val allMetrics = metrics.metrics()

    // If using custom quota callbacks or if multiple-levels of quotas are defined or
    // if this is a default quota update, traverse metrics to find all affected values.
    // Otherwise, update just the single matching one.
    val singleUpdate = quotaTypesEnabled match {
      case QuotaTypes.NoQuotas | QuotaTypes.ClientIdQuotaEnabled | QuotaTypes.UserQuotaEnabled | QuotaTypes.UserClientIdQuotaEnabled =>
        updatedQuotaEntity.nonEmpty
      case _ => false
    }
    if (singleUpdate) {
      val quotaEntity = updatedQuotaEntity.getOrElse(throw new IllegalStateException("Quota entity not specified"))

Review comment (suggestion):
let's say we have positive counts for QuotaTypes.UserClientIdQuotaEnabled and QuotaTypes.UserQuotaEnabled, when we print the `quotaTypesEnabled` we would expect to see `6` printed? perhaps out of scope for this change, but it would definitely be more friendly if the specific quotas enabled were printed instead.

Code after:
   *  @param shouldAdd   A boolean indicating whether to add or remove the quota entity.
   */
  private def updateQuotaTypes(quotaEntity: KafkaQuotaEntity, shouldAdd: Boolean): Unit = {

    if (quotaTypesEnabled == QuotaTypes.CustomQuotas) {
      // If custom quotas are enabled, we do not need to update quota types
      return
    }

    val isActive =  quotaCallback match {
      case callback: DefaultQuotaCallback => callback.getActiveQuotasEntities.contains(quotaEntity)
      case _ => true
    }

    val activeQuotaType = quotaEntity match {
      case KafkaQuotaEntity(Some(_), Some(_)) => QuotaTypes.UserClientIdQuotaEnabled
      case KafkaQuotaEntity(Some(_), None) => QuotaTypes.UserQuotaEnabled
      case KafkaQuotaEntity(None, Some(_)) => QuotaTypes.ClientIdQuotaEnabled
      case _ => QuotaTypes.NoQuotas
    }

    if (shouldAdd && !isActive) {
      activeQuotaEntities.compute(activeQuotaType, (_, currentValue) => if (currentValue == 0) 1 else currentValue + 1)
      quotaTypesEnabled |= activeQuotaType
    } else if (!shouldAdd && isActive) {
      activeQuotaEntities.compute(activeQuotaType, (_, currentValue) => if (currentValue <= 1) 0 else currentValue - 1)
      if (activeQuotaEntities.get(activeQuotaType) == 0) {
        quotaTypesEnabled &= ~activeQuotaType
      }
    }

    val quotaTypes = List(
      QuotaTypes.UserClientIdQuotaEnabled -> "UserClientIdQuota",
      QuotaTypes.ClientIdQuotaEnabled     -> "ClientIdQuota",
      QuotaTypes.UserQuotaEnabled         -> "UserQuota"
    )

    val activeEntities = quotaTypes.collect {
      case (k, name) if activeQuotaEntities.get(k) > 0 => name
    }.mkString(", ")
    info(s"Quota types enabled has been changed to $quotaTypesEnabled with active quota entities: [$activeEntities]")
  }

  /**
   * Updates metrics configs. This is invoked when quota configs are updated when partitions leaders change 
   * and custom callbacks that implement partition-based quotas have updated quotas.
   *
   * @param updatedQuotaEntity If set to one entity and quotas have only been enabled at one
   *    level, then an optimized update is performed with a single metric update. If None is provided,
   *    or if custom callbacks are used or if multi-level quotas have been enabled, all metric configs
   *    are checked and updated if required.