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

Code before:
   * Helper method to update quotaTypesEnabled which is a bitwise OR combination of the enabled quota types.
   * For example:
   *  - If UserQuotaEnabled = 2 and ClientIdQuotaEnabled = 1, then quotaTypesEnabled = 3 (2 | 1 = 3)
   *  - If UserClientIdQuotaEnabled = 4 and UserQuotaEnabled = 1, then quotaTypesEnabled =  (4 | 1 = 5)
   *  - If UserClientIdQuotaEnabled = 4 and ClientIdQuotaEnabled = 2, then quotaTypesEnabled = 6 (4 | 2 = 6)
   *  - If all three are enabled (1 | 2 | 4), then quotaTypesEnabled = 7
   */
  private def updateQuotaTypes(): Unit = {
    quotaTypesEnabled =  if (clientQuotaCallbackPlugin.isDefined) {
        QuotaTypes.CustomQuotas
      } else {
        QuotaTypes.NoQuotas
      }

    activeQuotaEntities.forEach {
        case KafkaQuotaEntity(Some(_), Some(_)) =>
          quotaTypesEnabled |= QuotaTypes.UserClientIdQuotaEnabled
        case KafkaQuotaEntity(Some(_), None) =>
          quotaTypesEnabled |= QuotaTypes.UserQuotaEnabled
        case KafkaQuotaEntity(None, Some(_)) =>
          quotaTypesEnabled |= QuotaTypes.ClientIdQuotaEnabled
        case _ => // Unexpected entity type
    }

    val activeEntities = if (activeQuotaEntities.isEmpty) "No active quota entities" else activeQuotaEntities.asScala.map(_.toString).mkString(", ")
    info(s"Quota types enabled has been changed 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.
   */
  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"))
      val user = quotaEntity.sanitizedUser

Review comment (bug):
Generating long string could be costly. Does it bring real value? Is this something that we could log in debug?

Code after:
   * This method ensures that the `quotaTypesEnabled` field reflects the active quota types based on the
   * current state of `activeQuotaEntities`.
   * For example:
   *  - If UserQuotaEnabled = 2 and ClientIdQuotaEnabled = 1, then quotaTypesEnabled = 3 (2 | 1 = 3)
   *  - If UserClientIdQuotaEnabled = 4 and UserQuotaEnabled = 1, then quotaTypesEnabled =  (4 | 1 = 5)
   *  - If UserClientIdQuotaEnabled = 4 and ClientIdQuotaEnabled = 2, then quotaTypesEnabled = 6 (4 | 2 = 6)
   *  - If all three are enabled (1 | 2 | 4), then quotaTypesEnabled = 7
   *
   *  @param quotaEntity The entity for which the quota is being updated, which can be a combination of user and client-id.
   *  @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]")
  }