From d1ffeb7fb7de0423c5c150f4f90296cb84a6b636 Mon Sep 17 00:00:00 2001 From: pablo Date: Sun, 6 Jul 2025 13:28:38 +0300 Subject: [PATCH] fix: changed ratio to actual capacity and PR comments --- tests/common/test_tokenbucket.nim | 66 ++++++++++++++++++++----- waku/common/rate_limit/token_bucket.nim | 16 ++++-- 2 files changed, 67 insertions(+), 15 deletions(-) diff --git a/tests/common/test_tokenbucket.nim b/tests/common/test_tokenbucket.nim index f0ce1fe1c..b7d1742f9 100644 --- a/tests/common/test_tokenbucket.nim +++ b/tests/common/test_tokenbucket.nim @@ -10,13 +10,8 @@ import testutils/unittests import chronos -import std/math import ../../waku/common/rate_limit/token_bucket -# Helper function for approximate float equality -proc equals(a, b: float, tolerance: float = 1e-6): bool = - abs(a - b) <= tolerance - suite "Token Bucket": test "TokenBucket Sync test - strict": var bucket = TokenBucket.newStrict(1000, 1.milliseconds) @@ -73,27 +68,76 @@ suite "Token Bucket": check bucket.tryConsume(15000, start + 1.milliseconds) == true - test "TokenBucket getAvailableCapacityRatio": + test "TokenBucket getAvailableCapacity strict": var bucket = TokenBucket.new(1000, 1.minutes, ReplenishMode.Strict) var reqTime = Moment.now() # Test full bucket capacity ratio - check equals(bucket.getAvailableCapacityRatio(reqTime), 1.0) # 1000/1000 = 1.0 + check bucket.getAvailableCapacity(reqTime) == (1000, 1000) # Consume some tokens and check ratio reqTime += 1.seconds check bucket.tryConsume(400, reqTime) == true - check equals(bucket.getAvailableCapacityRatio(reqTime), 0.6) # 600/1000 = 0.6 + check bucket.getAvailableCapacity(reqTime) == (600, 1000) # Consume more tokens reqTime += 1.seconds check bucket.tryConsume(300, reqTime) == true - check equals(bucket.getAvailableCapacityRatio(reqTime), 0.3) # 300/1000 = 0.3 + check bucket.getAvailableCapacity(reqTime) == (300, 1000) # Test when period has elapsed (should return 1.0) reqTime += 1.minutes - check equals(bucket.getAvailableCapacityRatio(reqTime), 1.0) # 1000/1000 = 1.0 + check bucket.getAvailableCapacity(reqTime) == (1000, 1000) # Test with empty bucket check bucket.tryConsume(1000, reqTime) == true - check equals(bucket.getAvailableCapacityRatio(reqTime), 0.0) # 0/1000 = 0.0 + check bucket.getAvailableCapacity(reqTime) == (0, 1000) + + test "TokenBucket getAvailableCapacity compensating": + var bucket = TokenBucket.new(1000, 1.minutes, ReplenishMode.Compensating) + var reqTime = Moment.now() + + # Test full bucket capacity + check bucket.getAvailableCapacity(reqTime) == (1000, 1000) + + # Consume some tokens and check available capacity + reqTime += 1.seconds + check bucket.tryConsume(400, reqTime) == true + check bucket.getAvailableCapacity(reqTime) == (600, 1000) + + # Consume more tokens + reqTime += 1.seconds + check bucket.tryConsume(300, reqTime) == true + check bucket.getAvailableCapacity(reqTime) == (300, 1000) + + # Test compensation when period has elapsed - should get compensation for unused capacity + # We used 700 tokens out of 1000 in 2 periods, so average usage was 35% per period + # Compensation should be added for the unused 65% capacity (up to 25% max) + reqTime += 1.minutes + let (availableBudget, maxCap) = bucket.getAvailableCapacity(reqTime) + check maxCap == 1000 + check availableBudget >= 1000 # Should have compensation + check availableBudget <= 1250 # But limited to 25% max compensation + + # Test with minimal usage - less consumption means less compensation + bucket = TokenBucket.new(1000, 1.minutes, ReplenishMode.Compensating) + reqTime = Moment.now() + check bucket.tryConsume(50, reqTime) == true # Use only 5% of capacity (950 remaining) + + # Move to next period - compensation based on remaining budget + # UsageAverage = 950/1000/1.0 = 0.95, so compensation = (1.0-0.95)*1000 = 50 + reqTime += 1.minutes + let (compensatedBudget, _) = bucket.getAvailableCapacity(reqTime) + check compensatedBudget == 1050 # 1000 + 50 compensation + + # Test with full usage - maximum compensation due to zero remaining budget + bucket = TokenBucket.new(1000, 1.minutes, ReplenishMode.Compensating) + reqTime = Moment.now() + check bucket.tryConsume(1000, reqTime) == true # Use full capacity (0 remaining) + + # Move to next period - maximum compensation since usage average is 0 + # UsageAverage = 0/1000/1.0 = 0.0, so compensation = (1.0-0.0)*1000 = 1000, capped at 250 + reqTime += 1.minutes + let (maxCompensationBudget, _) = bucket.getAvailableCapacity(reqTime) + check maxCompensationBudget == 1250 # 1000 + 250 max compensation + diff --git a/waku/common/rate_limit/token_bucket.nim b/waku/common/rate_limit/token_bucket.nim index b938ef039..6fe60bc80 100644 --- a/waku/common/rate_limit/token_bucket.nim +++ b/waku/common/rate_limit/token_bucket.nim @@ -109,11 +109,19 @@ proc update(bucket: TokenBucket, currentTime: Moment) = else: updateStrict(bucket, currentTime) -## Returns the available capacity ratio of the bucket: 0.0 (empty); 1.0 (full) -proc getAvailableCapacityRatio*(bucket: TokenBucket, currentTime: Moment): float = +## Returns the available capacity of the bucket: (budget, budgetCap) +proc getAvailableCapacity*(bucket: TokenBucket, currentTime: Moment = Moment.now()): tuple[budget: int, budgetCap: int] = if periodElapsed(bucket, currentTime): - return 1.0 - return bucket.budget.float / bucket.budgetCap.float + case bucket.replenishMode + of ReplenishMode.Strict: + return (bucket.budgetCap, bucket.budgetCap) + of ReplenishMode.Compensating: + let distance = bucket.periodDistance(currentTime) + let recentAvgUsage = bucket.getUsageAverageSince(distance) + let compensation = bucket.calcCompensation(recentAvgUsage) + let availableBudget = bucket.budgetCap + compensation + return (availableBudget, bucket.budgetCap) + return (bucket.budget, bucket.budgetCap) proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool = ## If `tokens` are available, consume them,