diff --git a/internal/rootcoord/quota_center.go b/internal/rootcoord/quota_center.go index 98e42f5e24676..511f465c65edb 100644 --- a/internal/rootcoord/quota_center.go +++ b/internal/rootcoord/quota_center.go @@ -1182,6 +1182,15 @@ func (q *QuotaCenter) calculateRates() error { 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 { @@ -1192,21 +1201,20 @@ func (q *QuotaCenter) resetAllCurrentRates() error { } 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)) } - if len(partitionIDs) == 0 { - q.rateLimiter.GetOrCreateCollectionLimiters(dbID, collectionID, - newParamLimiterFunc(internalpb.RateScope_Database, allOps), - newParamLimiterFuncWithLimitFunc(internalpb.RateScope_Collection, allOps, getCollectionLimitVal)) - } - } - if len(collections) == 0 { - q.rateLimiter.GetOrCreateDatabaseLimiters(dbID, newParamLimiterFunc(internalpb.RateScope_Database, allOps)) } } }