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

Code before:
   * for any of these levels.
   *
   * @param userEntity   user to override if quota applies to <user> or <user, client-id>
   * @param clientEntity sanitized client entity to override if quota applies to <client-id> or <user, client-id>
   * @param quota        custom quota to apply or None if quota override is being removed
   */
  def updateQuota(
    userEntity: Option[BaseUserEntity],
    clientEntity: Option[ClientQuotaEntity.ConfigEntity],
    quota: Option[Quota]
  ): Unit = {
    /*
     * 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 {
      val quotaEntity = KafkaQuotaEntity(userEntity, clientEntity)

      quota match {
        case Some(newQuota) =>
          quotaCallback.updateQuota(clientQuotaType, quotaEntity, newQuota.bound)
          if(!activeQuotaEntities.put(quotaEntity, true)){
            updateQuotaTypes()
          }
        case None =>
          quotaCallback.removeQuota(clientQuotaType, quotaEntity)
          if (activeQuotaEntities.remove(quotaEntity)){
            updateQuotaTypes()
          }
      }

      val updatedEntity = if (userEntity.contains(DefaultUserEntity) || clientEntity.contains(DefaultClientIdEntity))
        None // more than one entity may need updating, so `updateQuotaMetricConfigs` will go through all metrics
      else
        Some(quotaEntity)
      updateQuotaMetricConfigs(updatedEntity)

    } finally {
      lock.writeLock().unlock()
    }
  }

  /**
   * 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)

Review comment (suggestion):
Nit; do we ever insert a `false` value? If not, consider using a `HashSet` instead. 

Code after:
   * Overrides quotas for <user>, <client-id> or <user, client-id> or the dynamic defaults
   * for any of these levels.
   *
   * @param userEntity   user to override if quota applies to <user> or <user, client-id>
   * @param clientEntity sanitized client entity to override if quota applies to <client-id> or <user, client-id>
   * @param quota        custom quota to apply or None if quota override is being removed
   */
  def updateQuota(
    userEntity: Option[BaseUserEntity],
    clientEntity: Option[ClientQuotaEntity.ConfigEntity],
    quota: Option[Quota]
  ): Unit = {
    /*
     * 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 {
      val quotaEntity = KafkaQuotaEntity(userEntity, clientEntity)

      quota match {
        case Some(newQuota) =>
          updateQuotaTypes(quotaEntity, shouldAdd = true)
          quotaCallback.updateQuota(clientQuotaType, quotaEntity, newQuota.bound)

        case None =>
          updateQuotaTypes(quotaEntity, shouldAdd = false)
          quotaCallback.removeQuota(clientQuotaType, quotaEntity)
      }

      val updatedEntity = if (userEntity.contains(DefaultUserEntity) || clientEntity.contains(DefaultClientIdEntity))
        None // more than one entity may need updating, so `updateQuotaMetricConfigs` will go through all metrics
      else
        Some(quotaEntity)
      updateQuotaMetricConfigs(updatedEntity)

    } finally {
      lock.writeLock().unlock()
    }
  }

  /**
   * 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)