Repository: milvus-io/milvus
File: internal/rootcoord/quota_center.go (Go)

Code before:
	err = q.calculateWriteRates()
	if err != nil {
		log.Warn("QuotaCenter calculateWriteRates failed", zap.Error(err))
		return err
	}
	err = q.calculateReadRates()
	if err != nil {
		log.Warn("QuotaCenter calculateReadRates failed", zap.Error(err))
		return err
	}

	// log.Debug("QuotaCenter calculates rate done", zap.Any("rates", q.currentRates))
	return nil
}

func (q *QuotaCenter) resetAllCurrentRates() error {
	clusterLimiter := newParamLimiterFunc(internalpb.RateScope_Cluster, allOps)()
	q.rateLimiter = rlinternal.NewRateLimiterTree(clusterLimiter)

	enablePartitionRateLimit := false
	for rt := range getRateTypes(internalpb.RateScope_Partition, allOps) {
		r := quota.GetQuotaValue(internalpb.RateScope_Partition, rt, Params)
		if Limit(r) != Inf {
			enablePartitionRateLimit = true
		}
	}

	initLimiters := func(sourceCollections map[int64]map[int64][]int64) {
		for dbID, collections := range sourceCollections {
			for collectionID, partitionIDs := range collections {
				getCollectionLimitVal := func(rateType internalpb.RateType) Limit {
					limitVal, err := q.getCollectionMaxLimit(rateType, collectionID)
					if err != nil {
						return Limit(quota.GetQuotaValue(internalpb.RateScope_Collection, rateType, Params))
					}
					return limitVal
				}
				q.rateLimiter.GetOrCreateDatabaseLimiters(dbID, newParamLimiterFunc(internalpb.RateScope_Database, allOps))
				q.rateLimiter.GetOrCreateCollectionLimiters(dbID, collectionID,
					newParamLimiterFunc(internalpb.RateScope_Database, allOps),
					newParamLimiterFuncWithLimitFunc(internalpb.RateScope_Collection, allOps, getCollectionLimitVal))

				if !enablePartitionRateLimit {
					continue
				}
				for _, partitionID := range partitionIDs {
					q.rateLimiter.GetOrCreatePartitionLimiters(dbID, collectionID, partitionID,
						newParamLimiterFunc(internalpb.RateScope_Database, allOps),
						newParamLimiterFuncWithLimitFunc(internalpb.RateScope_Collection, allOps, getCollectionLimitVal),
						newParamLimiterFunc(internalpb.RateScope_Partition, allOps))
				}

Review comment (suggestion):
It implies `Inf` are the same as `nil` for partition limiter. And this context need to be maintained by `rate_limiter_tree` instead of this function. 

Suggest to create a new function in `rate_limiter_tree` to skip partition limiter's construction.

Code after:
	err = q.calculateWriteRates()
	if err != nil {
		log.Warn("QuotaCenter calculateWriteRates failed", zap.Error(err))
		return err
	}
	err = q.calculateReadRates()
	if err != nil {
		log.Warn("QuotaCenter calculateReadRates failed", zap.Error(err))
		return err
	}

	// log.Debug("QuotaCenter calculates rate done", zap.Any("rates", q.currentRates))
	return nil
}

func (q *QuotaCenter) resetAllCurrentRates() error {
	clusterLimiter := newParamLimiterFunc(internalpb.RateScope_Cluster, allOps)()
	q.rateLimiter = rlinternal.NewRateLimiterTree(clusterLimiter)

	enablePartitionRateLimit := false
	for rt := range getRateTypes(internalpb.RateScope_Partition, allOps) {
		r := quota.GetQuotaValue(internalpb.RateScope_Partition, rt, Params)
		if Limit(r) != Inf {
			enablePartitionRateLimit = true
		}
	}

	initLimiters := func(sourceCollections map[int64]map[int64][]int64) {
		for dbID, collections := range sourceCollections {
			for collectionID, partitionIDs := range collections {
				getCollectionLimitVal := func(rateType internalpb.RateType) Limit {
					limitVal, err := q.getCollectionMaxLimit(rateType, collectionID)
					if err != nil {
						return Limit(quota.GetQuotaValue(internalpb.RateScope_Collection, rateType, Params))
					}
					return limitVal
				}
				q.rateLimiter.GetOrCreateCollectionLimiters(dbID, collectionID,
					newParamLimiterFunc(internalpb.RateScope_Database, allOps),
					newParamLimiterFuncWithLimitFunc(internalpb.RateScope_Collection, allOps, getCollectionLimitVal))

				if !enablePartitionRateLimit {
					continue
				}
				for _, partitionID := range partitionIDs {
					q.rateLimiter.GetOrCreatePartitionLimiters(dbID, collectionID, partitionID,
						newParamLimiterFunc(internalpb.RateScope_Database, allOps),
						newParamLimiterFuncWithLimitFunc(internalpb.RateScope_Collection, allOps, getCollectionLimitVal),
						newParamLimiterFunc(internalpb.RateScope_Partition, allOps))
				}
			}