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

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

Review comment (suggestion):
`quotaTypesEnabled` is a volatile variable. This suggests that we can access it anytime. In this method, we update it in a non atomic manner. I wonder if this could have any undesired impact. It would be better to use a local variable and to update `quotaTypesEnabled` at the end.

Code after:
          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)
   *  - 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) {