feat(resourcemanager): add trickle-time enforcement to service limiter (#9684)

ref #9296, close #9668

Enhance `serviceLimiter.applyServiceLimit` to return both limited tokens and a minimum trickle time, ensuring client-side rate does not exceed service limit.

Signed-off-by: JmPotato <[email protected]>

Author: JmPotato

pkg/mcs/resourcemanager/server/manager.go

@@ -530,7 +530,7 @@ func (m *Manager) backgroundMetricsFlush(ctx context.Context) {
 				if err != nil {
 					continue
 				}
-				setOrRemoveServiceLimitMetrics(keyspaceName, krgm.getServiceLimiter().GetServiceLimit())
+				setOrRemoveServiceLimitMetrics(keyspaceName, krgm.getServiceLimiter().getServiceLimit())
 
 				for _, group := range krgm.getResourceGroupList(true, true) {
 					groupName := group.Name

pkg/mcs/resourcemanager/server/resource_group.go

@@ -273,14 +273,17 @@ func (rg *ResourceGroup) RequestRU(
 	}
 	// Then, try to apply the service limit.
 	grantedTokens := tb.GetTokens()
-	limitedTokens := sl.applyServiceLimit(now, grantedTokens)
+	limitedTokens, minTrickleTimeMs := sl.applyServiceLimit(now, grantedTokens)
 	if limitedTokens < grantedTokens {
 		tb.Tokens = limitedTokens
 		// Retain the unused tokens for the later requests if it has a burst limit.
 		if rg.getBurstLimitLocked() > 0 {
 			rg.RUSettings.RU.reservedServiceTokens += grantedTokens - limitedTokens
 		}
 	}
+	if trickleTimeMs < minTrickleTimeMs {
+		trickleTimeMs = minTrickleTimeMs
+	}
 	return &rmpb.GrantedRUTokenBucket{GrantedTokens: tb, TrickleTimeMs: trickleTimeMs}
 }
 

pkg/mcs/resourcemanager/server/service_limit.go

@@ -91,13 +91,16 @@ func (krl *serviceLimiter) setServiceLimit(newServiceLimit float64) {
 	}
 }
 
-// GetServiceLimit return the service limit value of this keyspace.
-func (krl *serviceLimiter) GetServiceLimit() float64 {
+func (krl *serviceLimiter) getServiceLimit() float64 {
 	if krl == nil {
 		return 0.0
 	}
 	krl.RLock()
 	defer krl.RUnlock()
+	return krl.getServiceLimitLocked()
+}
+
+func (krl *serviceLimiter) getServiceLimitLocked() float64 {
 	return krl.ServiceLimit
 }
 
@@ -130,35 +133,35 @@ func (krl *serviceLimiter) refillTokensLocked(now time.Time) {
 func (krl *serviceLimiter) applyServiceLimit(
 	now time.Time,
 	requestedTokens float64,
-) (limitedTokens float64) {
+) (limitedTokens float64, minTrickleTimeMs int64) {
 	if krl == nil {
-		return requestedTokens
+		return requestedTokens, 0
 	}
 	krl.Lock()
 	defer krl.Unlock()
 
 	// No limit configured, allow all tokens.
 	if krl.ServiceLimit <= 0 {
-		return requestedTokens
+		return requestedTokens, 0
 	}
 
 	// Refill first to ensure the available tokens is up to date.
 	krl.refillTokensLocked(now)
 
 	// If the requested tokens is less than the available tokens, grant all tokens.
 	if requestedTokens <= krl.AvailableTokens {
+		limitedTokens = requestedTokens
 		krl.AvailableTokens -= requestedTokens
-		return requestedTokens
-	}
-
-	// If the requested tokens is greater than the available tokens, grant all available tokens.
-	if requestedTokens > krl.AvailableTokens {
+	} else {
+		// If the requested tokens is greater than the available tokens, grant all available tokens.
 		limitedTokens = math.Max(0, krl.AvailableTokens)
-		// TODO: allow the loan to decrease the allocation at a smooth rate.
 		krl.AvailableTokens = 0
 	}
 
-	return limitedTokens
+	// Calculate a minimum trickle time to ensure the granted tokens' rate in client won't exceed the service limit.
+	minTrickleTimeMs = int64(math.Round(limitedTokens * 1000.0 / krl.getServiceLimitLocked()))
+
+	return limitedTokens, minTrickleTimeMs
 }
 
 // Clone returns a copy of the service limiter.

pkg/mcs/resourcemanager/server/service_limit_test.go

@@ -135,38 +135,47 @@ func TestApplyServiceLimit(t *testing.T) {
 
 	// Test with nil limiter
 	var limiter *serviceLimiter
-	tokens := limiter.applyServiceLimit(time.Now(), 50.0)
+	now := time.Now()
+	tokens, minTrickleTimeMs := limiter.applyServiceLimit(now, 50.0)
 	re.Equal(50.0, tokens)
+	re.Zero(minTrickleTimeMs)
 
 	// Test with zero service limit (no limit)
 	limiter = newServiceLimiter(constant.NullKeyspaceID, 0.0, nil)
-	now := time.Now()
-	tokens = limiter.applyServiceLimit(now, 50.0)
+	now = time.Now()
+	tokens, minTrickleTimeMs = limiter.applyServiceLimit(now, 50.0)
 	re.Equal(50.0, tokens)
+	re.Zero(minTrickleTimeMs)
 
 	// Test request within available tokens (need to set available tokens first)
 	limiter = newServiceLimiter(constant.NullKeyspaceID, 100.0, nil)
+	now = time.Now()
 	limiter.AvailableTokens = 100.0 // Manually set available tokens
 	limiter.LastUpdate = now
-	tokens = limiter.applyServiceLimit(now, 50.0)
+	tokens, minTrickleTimeMs = limiter.applyServiceLimit(now, 50.0)
 	re.Equal(50.0, tokens)
 	re.Equal(50.0, limiter.AvailableTokens) // 100 - 50 = 50
+	re.Equal(int64(500), minTrickleTimeMs)  // 50/100 * 1000 = 500ms
 
 	// Test request exactly equal to available tokens
 	limiter = newServiceLimiter(constant.NullKeyspaceID, 100.0, nil)
+	now = time.Now()
 	limiter.AvailableTokens = 100.0 // Manually set available tokens
 	limiter.LastUpdate = now
-	tokens = limiter.applyServiceLimit(now, 100.0)
+	tokens, minTrickleTimeMs = limiter.applyServiceLimit(now, 100.0)
 	re.Equal(100.0, tokens)
 	re.Equal(0.0, limiter.AvailableTokens)
+	re.Equal(int64(1000), minTrickleTimeMs) // 100/100 * 1000 = 1000ms
 
 	// Test request exceeding available tokens
 	limiter = newServiceLimiter(constant.NullKeyspaceID, 100.0, nil)
+	now = time.Now()
 	limiter.LastUpdate = now
 	limiter.AvailableTokens = 30.0
-	tokens = limiter.applyServiceLimit(now, 80.0)
+	tokens, minTrickleTimeMs = limiter.applyServiceLimit(now, 80.0)
 	re.Equal(30.0, tokens) // Only available tokens granted
 	re.Equal(0.0, limiter.AvailableTokens)
+	re.Equal(int64(300), minTrickleTimeMs) // 30/100 * 1000 = 300ms
 }
 
 func TestApplyServiceLimitWithRefill(t *testing.T) {
@@ -180,21 +189,23 @@ func TestApplyServiceLimitWithRefill(t *testing.T) {
 
 	// Request after 1 second should trigger refill first
 	futureTime := baseTime.Add(time.Second)
-	tokens := limiter.applyServiceLimit(futureTime, 50.0)
+	tokens, minTrickleTimeMs := limiter.applyServiceLimit(futureTime, 50.0)
 	re.Equal(50.0, tokens)
 	re.Equal(70.0, limiter.AvailableTokens) // 20 + 100 - 50 = 70
 	re.Equal(futureTime, limiter.LastUpdate)
+	re.Equal(int64(500), minTrickleTimeMs) // 50/100 * 1000 = 500ms
 
 	// Test partial refill scenario
 	limiter = newServiceLimiter(constant.NullKeyspaceID, 100.0, nil)
 	limiter.LastUpdate = baseTime
 	limiter.AvailableTokens = 10.0
 
 	halfSecondLater := baseTime.Add(500 * time.Millisecond)
-	tokens = limiter.applyServiceLimit(halfSecondLater, 80.0)
+	tokens, minTrickleTimeMs = limiter.applyServiceLimit(halfSecondLater, 80.0)
 	// After refill: 10 + 100*0.5 = 60 tokens available
 	re.Equal(60.0, tokens)
 	re.Equal(0.0, limiter.AvailableTokens)
+	re.Equal(int64(600), minTrickleTimeMs) // 60/100 * 1000 = 600ms
 }
 
 func TestServiceLimiterEdgeCases(t *testing.T) {
@@ -205,32 +216,37 @@ func TestServiceLimiterEdgeCases(t *testing.T) {
 	limiter.AvailableTokens = 0.1   // Manually set available tokens
 	limiter.LastUpdate = time.Now() // Set LastUpdate to current time to avoid refill
 	now := time.Now()
-	tokens := limiter.applyServiceLimit(now, 1.0)
-	re.InDelta(0.1, tokens, 0.001) // Use InDelta to handle floating point precision
+	tokens, minTrickleTimeMs := limiter.applyServiceLimit(now, 1.0)
+	re.InDelta(0.1, tokens, 0.001)          // Use InDelta to handle floating point precision
+	re.Equal(int64(1000), minTrickleTimeMs) // 0.1/0.1 * 1000 = 1000ms
 
 	// Test with very large service limit
 	limiter = newServiceLimiter(constant.NullKeyspaceID, 1000000.0, nil)
 	limiter.AvailableTokens = 1000000.0 // Manually set available tokens
-	tokens = limiter.applyServiceLimit(now, 500000.0)
+	tokens, minTrickleTimeMs = limiter.applyServiceLimit(now, 500000.0)
 	re.Equal(500000.0, tokens)
+	re.Equal(int64(500), minTrickleTimeMs) // 500000/1000000 * 1000 = 500ms
 
 	// Test with zero requested tokens
 	limiter = newServiceLimiter(constant.NullKeyspaceID, 100.0, nil)
 	limiter.AvailableTokens = 100.0 // Manually set available tokens
-	tokens = limiter.applyServiceLimit(now, 0.0)
+	tokens, minTrickleTimeMs = limiter.applyServiceLimit(now, 0.0)
 	re.Equal(0.0, tokens)
 	re.Equal(100.0, limiter.AvailableTokens) // Should remain unchanged
+	re.Zero(minTrickleTimeMs)
 
 	// Test with fractional tokens
 	limiter = newServiceLimiter(constant.NullKeyspaceID, 10.5, nil)
 	limiter.LastUpdate = now
 	limiter.AvailableTokens = 5.25
-	tokens = limiter.applyServiceLimit(now, 7.75)
+	tokens, minTrickleTimeMs = limiter.applyServiceLimit(now, 7.75)
 	re.Equal(5.25, tokens)
+	re.Equal(int64(500), minTrickleTimeMs) // 5.25/10.5 * 1000 = 500ms
 	// Test apply with 0 available ru
-	tokens = limiter.applyServiceLimit(now, 5)
+	tokens, minTrickleTimeMs = limiter.applyServiceLimit(now, 5)
 	re.Equal(0.0, tokens)
 	re.Equal(0.0, limiter.AvailableTokens)
+	re.Zero(minTrickleTimeMs)
 }
 
 func TestSetServiceLimit(t *testing.T) {

pkg/mcs/resourcemanager/server/token_buckets.go

@@ -438,7 +438,7 @@ func (gtb *GroupTokenBucket) inspectAnomalies(
 ) bool {
 	var errMsg string
 	// Verify whether the allocated token is invalid, such as negative values, math.Inf, or math.NaN.
-	if tb.Tokens <= 0 || math.IsInf(tb.Tokens, 0) || math.IsNaN(tb.Tokens) {
+	if tb.Tokens < 0 || math.IsInf(tb.Tokens, 0) || math.IsNaN(tb.Tokens) {
 		errMsg = "assigned token is invalid"
 	}
 	// Verify whether the state of the slot is abnormal.