enhance: Skip creating partition rate limiters when not enable

Signed-off-by: bigsheeper <[email protected]>
bigsheeper · Dec 27, 2024 · 4f9b170 · 4f9b170
1 parent d433b75
commit 4f9b170
Showing 1 changed file with 16 additions and 8 deletions.
diff --git 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))
 			}
 		}
 	}