From 8912fc20e9b5d2ee1fb2d48f0ecb5ef56c6d041c Mon Sep 17 00:00:00 2001 From: pablo Date: Wed, 16 Jul 2025 12:36:11 +0300 Subject: [PATCH] feat: allow initialization with a non-default state --- tests/common/test_tokenbucket.nim | 45 +++++++++++++++---------- waku/common/rate_limit/token_bucket.nim | 21 +++++++----- 2 files changed, 41 insertions(+), 25 deletions(-) diff --git a/tests/common/test_tokenbucket.nim b/tests/common/test_tokenbucket.nim index 5d6b150a8..12a02f905 100644 --- a/tests/common/test_tokenbucket.nim +++ b/tests/common/test_tokenbucket.nim @@ -69,52 +69,63 @@ suite "Token Bucket": check bucket.tryConsume(15000, start + 1.milliseconds) == true test "TokenBucket getAvailableCapacity strict": - var bucket = TokenBucket.new(1000, 1.minutes, ReplenishMode.Strict) - var reqTime = Moment.now() + let now = Moment.now() + var reqTime = now + var bucket = + TokenBucket.new(1000, 1.minutes, ReplenishMode.Strict, lastTimeFull = now) # Test full bucket capacity ratio - check bucket.getAvailableCapacity(reqTime) == (1000, 1000) + let lastTimeFull_1 = reqTime + check bucket.getAvailableCapacity(reqTime) == (1000, 1000, lastTimeFull_1) - # Consume some tokens and check ratio + # # Consume some tokens and check ratio reqTime += 1.seconds + let lastTimeFull_2 = reqTime + check bucket.tryConsume(400, reqTime) == true - check bucket.getAvailableCapacity(reqTime) == (600, 1000) + check bucket.getAvailableCapacity(reqTime) == (600, 1000, lastTimeFull_2) # Consume more tokens reqTime += 1.seconds check bucket.tryConsume(300, reqTime) == true - check bucket.getAvailableCapacity(reqTime) == (300, 1000) + check bucket.getAvailableCapacity(reqTime) == (300, 1000, lastTimeFull_2) - # Test when period has elapsed (should return 1.0) + # Test when period has elapsed (should return 1.0) reqTime += 1.minutes - check bucket.getAvailableCapacity(reqTime) == (1000, 1000) + check bucket.getAvailableCapacity(reqTime) == (1000, 1000, lastTimeFull_2) + let lastTimeFull_3 = reqTime # Test with empty bucket check bucket.tryConsume(1000, reqTime) == true - check bucket.getAvailableCapacity(reqTime) == (0, 1000) + check bucket.getAvailableCapacity(reqTime) == (0, 1000, lastTimeFull_3) test "TokenBucket getAvailableCapacity compensating": - var bucket = TokenBucket.new(1000, 1.minutes, ReplenishMode.Compensating) - var reqTime = Moment.now() + let now = Moment.now() + var reqTime = now + var bucket = + TokenBucket.new(1000, 1.minutes, ReplenishMode.Compensating, lastTimeFull = now) + + let lastTimeFull_1 = reqTime # Test full bucket capacity - check bucket.getAvailableCapacity(reqTime) == (1000, 1000) + check bucket.getAvailableCapacity(reqTime) == (1000, 1000, lastTimeFull_1) # Consume some tokens and check available capacity reqTime += 1.seconds + let lastTimeFull_2 = reqTime check bucket.tryConsume(400, reqTime) == true - check bucket.getAvailableCapacity(reqTime) == (600, 1000) + check bucket.getAvailableCapacity(reqTime) == (600, 1000, lastTimeFull_2) # Consume more tokens reqTime += 1.seconds check bucket.tryConsume(300, reqTime) == true - check bucket.getAvailableCapacity(reqTime) == (300, 1000) + check bucket.getAvailableCapacity(reqTime) == (300, 1000, lastTimeFull_2) # 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) + let (availableBudget, maxCap, _) = bucket.getAvailableCapacity(reqTime) check maxCap == 1000 check availableBudget >= 1000 # Should have compensation check availableBudget <= 1250 # But limited to 25% max compensation @@ -128,7 +139,7 @@ suite "Token Bucket": # 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) + let (compensatedBudget, _, _) = bucket.getAvailableCapacity(reqTime) check compensatedBudget == 1050 # 1000 + 50 compensation # Test with full usage - maximum compensation due to zero remaining budget @@ -139,5 +150,5 @@ suite "Token Bucket": # 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) + 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 9dcb2e7d1..f52cc8a97 100644 --- a/waku/common/rate_limit/token_bucket.nim +++ b/waku/common/rate_limit/token_bucket.nim @@ -111,20 +111,20 @@ proc update(bucket: TokenBucket, currentTime: Moment) = proc getAvailableCapacity*( bucket: TokenBucket, currentTime: Moment = Moment.now() -): tuple[budget: int, budgetCap: int] = +): tuple[budget: int, budgetCap: int, lastTimeFull: Moment] = ## Returns the available capacity of the bucket: (budget, budgetCap) if periodElapsed(bucket, currentTime): case bucket.replenishMode of ReplenishMode.Strict: - return (bucket.budgetCap, bucket.budgetCap) + return (bucket.budgetCap, bucket.budgetCap, bucket.lastTimeFull) 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) + return (availableBudget, bucket.budgetCap, bucket.lastTimeFull) + return (bucket.budget, bucket.budgetCap, bucket.lastTimeFull) proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool = ## If `tokens` are available, consume them, @@ -155,26 +155,31 @@ proc new*( budgetCap: int, fillDuration: Duration = 1.seconds, mode: ReplenishMode = ReplenishMode.Compensating, + budget: int = -1, # -1 means "use budgetCap" + lastTimeFull: Moment = Moment.now(), ): T = assert not isZero(fillDuration) assert budgetCap != 0 + assert lastTimeFull <= Moment.now() + let actualBudget = if budget == -1: budgetCap else: budget + assert actualBudget >= 0 and actualBudget <= budgetCap ## Create different mode TokenBucket case mode of ReplenishMode.Strict: return T( - budget: budgetCap, + budget: actualBudget, budgetCap: budgetCap, fillDuration: fillDuration, - lastTimeFull: Moment.now(), + lastTimeFull: lastTimeFull, replenishMode: mode, ) of ReplenishMode.Compensating: T( - budget: budgetCap, + budget: actualBudget, budgetCap: budgetCap, fillDuration: fillDuration, - lastTimeFull: Moment.now(), + lastTimeFull: lastTimeFull, replenishMode: mode, maxCompensation: budgetCap.float * BUDGET_COMPENSATION_LIMIT_PERCENT, )