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

Code before:
    }
  }

  /**
   * Updates `quotaTypesEnabled` by performing a bitwise OR operation to combine the enabled quota types.
   * 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 = {

    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
    }

    val isActive = quotaCallback.getActiveQuotasEntities.contains(quotaEntity)
    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.

Review comment (suggestion):
Since we cannot put `getActiveQuotasEntities` in the interface, we can instead do:
```
val isActive = quotaCallback match {
      case callback: DefaultQuotaCallback => callback.getActiveQuotasEntities.contains(quotaEntity)
      case _ => true
    }
```

Code after:
    }
  }

  /**
   * Updates `quotaTypesEnabled` by performing a bitwise OR operation to combine the enabled quota types.
   * 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"
    )