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

Code before:
   *
   * @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)

      val quotaTypes = (userEntity.nonEmpty, clientEntity.nonEmpty) match {
        case (true, true) => QuotaTypes.UserClientIdQuotaEnabled
        case (true, false) => QuotaTypes.UserQuotaEnabled
        case (false, true) => QuotaTypes.ClientIdQuotaEnabled
        case (false, false) => QuotaTypes.NoQuotas
      }
      quota match {
        case Some(newQuota) =>
          quotaCallback.updateQuota(clientQuotaType, quotaEntity, newQuota.bound)
          updateQuotaTypes(quotaTypes, increment = true)
        case None =>
          quotaCallback.removeQuota(clientQuotaType, quotaEntity) // change here.
          updateQuotaTypes(quotaTypes, increment = false)
      }

      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 quota types counts and quotaTypesEnabled flag.
   * @param quotaTypeKey The QuotaTypes constant (e.g., QuotaTypes.UserClientIdQuotaEnabled)
   * @param increment True to increment count, false to decrement

Review comment (suggestion):
just want to clarify my understanding of NoQuotas - if this method were called with a non-empty `quota` but empty `userEntity` and empty `clientEntity`, would that essentially be a no-op? assuming that combination is not expected/possible

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)