From a5ff6a1fbd8d031d349c46550157b6887ba43e0d Mon Sep 17 00:00:00 2001 From: pablo Date: Thu, 3 Jul 2025 12:10:15 +0300 Subject: [PATCH 01/10] feat: expose available needed by chatsdk --- tests/common/test_tokenbucket.nim | 30 +++++++++++++++++++++++++ waku/common/rate_limit/token_bucket.nim | 8 +++++++ 2 files changed, 38 insertions(+) diff --git a/tests/common/test_tokenbucket.nim b/tests/common/test_tokenbucket.nim index 5bc1a0583..413558e24 100644 --- a/tests/common/test_tokenbucket.nim +++ b/tests/common/test_tokenbucket.nim @@ -10,8 +10,13 @@ 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) @@ -67,3 +72,28 @@ suite "Token Bucket": check bucket.tryConsume(1, start) == false check bucket.tryConsume(15000, start + 1.milliseconds) == true + + test "TokenBucket getAvailableCapacityRatio": + 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 + + # 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 + + # Consume more tokens + reqTime += 1.seconds + check bucket.tryConsume(300, reqTime) == true + check equals(bucket.getAvailableCapacityRatio(reqTime), 0.3) # 300/1000 = 0.3 + + # Test when period has elapsed (should return 1.0) + reqTime += 1.minutes + check equals(bucket.getAvailableCapacityRatio(reqTime), 1.0) # 1000/1000 = 1.0 + + # Test with empty bucket + check bucket.tryConsume(1000, reqTime) == true + check equals(bucket.getAvailableCapacityRatio(reqTime), 0.0) # 0/1000 = 0.0 diff --git a/waku/common/rate_limit/token_bucket.nim b/waku/common/rate_limit/token_bucket.nim index 799817ebd..265d8847a 100644 --- a/waku/common/rate_limit/token_bucket.nim +++ b/waku/common/rate_limit/token_bucket.nim @@ -109,6 +109,14 @@ proc update(bucket: TokenBucket, currentTime: Moment) = else: updateStrict(bucket, currentTime) +## Returns the available capacity ratio of the bucket. +## This is a ratio of the bucket capacity that is currently available. +## It is a float number between 0.0 (empty) and 1.0 (full). +proc getAvailableCapacityRatio*(bucket: TokenBucket, currentTime: Moment): float = + if periodElapsed(bucket, currentTime): + return 1.0 + return bucket.budget.float / bucket.budgetCap.float + proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool = ## If `tokens` are available, consume them, ## Otherwhise, return false. From 0f3dfc82c6f4710e91ba263dea95745509b5fb42 Mon Sep 17 00:00:00 2001 From: pablo Date: Thu, 3 Jul 2025 12:12:20 +0300 Subject: [PATCH 02/10] feat: cleanup --- waku/common/rate_limit/token_bucket.nim | 1 - 1 file changed, 1 deletion(-) diff --git a/waku/common/rate_limit/token_bucket.nim b/waku/common/rate_limit/token_bucket.nim index 265d8847a..27e28f323 100644 --- a/waku/common/rate_limit/token_bucket.nim +++ b/waku/common/rate_limit/token_bucket.nim @@ -110,7 +110,6 @@ proc update(bucket: TokenBucket, currentTime: Moment) = updateStrict(bucket, currentTime) ## Returns the available capacity ratio of the bucket. -## This is a ratio of the bucket capacity that is currently available. ## It is a float number between 0.0 (empty) and 1.0 (full). proc getAvailableCapacityRatio*(bucket: TokenBucket, currentTime: Moment): float = if periodElapsed(bucket, currentTime): From ac5559fb5f84a2ae7324f4bbc45158505e2d17ed Mon Sep 17 00:00:00 2001 From: pablo Date: Thu, 3 Jul 2025 12:12:58 +0300 Subject: [PATCH 03/10] feat: cleanup --- waku/common/rate_limit/token_bucket.nim | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/waku/common/rate_limit/token_bucket.nim b/waku/common/rate_limit/token_bucket.nim index 27e28f323..a545c62fa 100644 --- a/waku/common/rate_limit/token_bucket.nim +++ b/waku/common/rate_limit/token_bucket.nim @@ -110,7 +110,7 @@ proc update(bucket: TokenBucket, currentTime: Moment) = updateStrict(bucket, currentTime) ## Returns the available capacity ratio of the bucket. -## It is a float number between 0.0 (empty) and 1.0 (full). +## as a float number between 0.0 (empty) and 1.0 (full). proc getAvailableCapacityRatio*(bucket: TokenBucket, currentTime: Moment): float = if periodElapsed(bucket, currentTime): return 1.0 From b29882638176dd2d49f150196780f9ec024f9ed5 Mon Sep 17 00:00:00 2001 From: pablo Date: Thu, 3 Jul 2025 12:14:13 +0300 Subject: [PATCH 04/10] feat: cleanup --- waku/common/rate_limit/token_bucket.nim | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/waku/common/rate_limit/token_bucket.nim b/waku/common/rate_limit/token_bucket.nim index a545c62fa..b938ef039 100644 --- a/waku/common/rate_limit/token_bucket.nim +++ b/waku/common/rate_limit/token_bucket.nim @@ -109,8 +109,7 @@ proc update(bucket: TokenBucket, currentTime: Moment) = else: updateStrict(bucket, currentTime) -## Returns the available capacity ratio of the bucket. -## as a float number between 0.0 (empty) and 1.0 (full). +## Returns the available capacity ratio of the bucket: 0.0 (empty); 1.0 (full) proc getAvailableCapacityRatio*(bucket: TokenBucket, currentTime: Moment): float = if periodElapsed(bucket, currentTime): return 1.0 From 91cf9c2703ba5b978d96e5d97a33a5ccdd174bd4 Mon Sep 17 00:00:00 2001 From: pablo Date: Thu, 3 Jul 2025 16:58:14 +0300 Subject: [PATCH 05/10] fix: lint --- tests/common/test_tokenbucket.nim | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/common/test_tokenbucket.nim b/tests/common/test_tokenbucket.nim index 413558e24..f0ce1fe1c 100644 --- a/tests/common/test_tokenbucket.nim +++ b/tests/common/test_tokenbucket.nim @@ -76,24 +76,24 @@ suite "Token Bucket": test "TokenBucket getAvailableCapacityRatio": 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 - + # 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 equals(bucket.getAvailableCapacityRatio(reqTime), 0.6) # 600/1000 = 0.6 + # Consume more tokens reqTime += 1.seconds check bucket.tryConsume(300, reqTime) == true - check equals(bucket.getAvailableCapacityRatio(reqTime), 0.3) # 300/1000 = 0.3 - + check equals(bucket.getAvailableCapacityRatio(reqTime), 0.3) # 300/1000 = 0.3 + # Test when period has elapsed (should return 1.0) reqTime += 1.minutes check equals(bucket.getAvailableCapacityRatio(reqTime), 1.0) # 1000/1000 = 1.0 - + # Test with empty bucket check bucket.tryConsume(1000, reqTime) == true - check equals(bucket.getAvailableCapacityRatio(reqTime), 0.0) # 0/1000 = 0.0 + check equals(bucket.getAvailableCapacityRatio(reqTime), 0.0) # 0/1000 = 0.0 From d1ffeb7fb7de0423c5c150f4f90296cb84a6b636 Mon Sep 17 00:00:00 2001 From: pablo Date: Sun, 6 Jul 2025 13:28:38 +0300 Subject: [PATCH 06/10] 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, From 340b82f7b60ce8c27e830eb00aad7a0779d7609d Mon Sep 17 00:00:00 2001 From: pablo Date: Mon, 7 Jul 2025 09:00:41 +0300 Subject: [PATCH 07/10] fix: lint --- tests/common/test_tokenbucket.nim | 18 +++++++++--------- waku/common/rate_limit/token_bucket.nim | 4 +++- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/tests/common/test_tokenbucket.nim b/tests/common/test_tokenbucket.nim index b7d1742f9..5d6b150a8 100644 --- a/tests/common/test_tokenbucket.nim +++ b/tests/common/test_tokenbucket.nim @@ -116,28 +116,28 @@ suite "Token Bucket": 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 + 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) - + 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 + 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) - + 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 - + 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 6fe60bc80..bb4baad99 100644 --- a/waku/common/rate_limit/token_bucket.nim +++ b/waku/common/rate_limit/token_bucket.nim @@ -110,7 +110,9 @@ proc update(bucket: TokenBucket, currentTime: Moment) = updateStrict(bucket, currentTime) ## Returns the available capacity of the bucket: (budget, budgetCap) -proc getAvailableCapacity*(bucket: TokenBucket, currentTime: Moment = Moment.now()): tuple[budget: int, budgetCap: int] = +proc getAvailableCapacity*( + bucket: TokenBucket, currentTime: Moment = Moment.now() +): tuple[budget: int, budgetCap: int] = if periodElapsed(bucket, currentTime): case bucket.replenishMode of ReplenishMode.Strict: From 7c0701ef83f5342852ae0f3d97e602e7082ae66c Mon Sep 17 00:00:00 2001 From: pablo Date: Tue, 8 Jul 2025 10:50:41 +0300 Subject: [PATCH 08/10] fix: comment --- waku/common/rate_limit/token_bucket.nim | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/waku/common/rate_limit/token_bucket.nim b/waku/common/rate_limit/token_bucket.nim index bb4baad99..9dcb2e7d1 100644 --- a/waku/common/rate_limit/token_bucket.nim +++ b/waku/common/rate_limit/token_bucket.nim @@ -109,10 +109,11 @@ proc update(bucket: TokenBucket, currentTime: Moment) = else: updateStrict(bucket, currentTime) -## Returns the available capacity of the bucket: (budget, budgetCap) proc getAvailableCapacity*( bucket: TokenBucket, currentTime: Moment = Moment.now() ): tuple[budget: int, budgetCap: int] = + ## Returns the available capacity of the bucket: (budget, budgetCap) + if periodElapsed(bucket, currentTime): case bucket.replenishMode of ReplenishMode.Strict: From 8912fc20e9b5d2ee1fb2d48f0ecb5ef56c6d041c Mon Sep 17 00:00:00 2001 From: pablo Date: Wed, 16 Jul 2025 12:36:11 +0300 Subject: [PATCH 09/10] 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, ) From 4e49dc26d2f64cb16d617fda4e29b8393344f2e5 Mon Sep 17 00:00:00 2001 From: pablo Date: Thu, 17 Jul 2025 08:57:48 +0300 Subject: [PATCH 10/10] feat: add unit tests --- tests/common/test_tokenbucket.nim | 97 +++++++++++++++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/tests/common/test_tokenbucket.nim b/tests/common/test_tokenbucket.nim index 12a02f905..786920ee6 100644 --- a/tests/common/test_tokenbucket.nim +++ b/tests/common/test_tokenbucket.nim @@ -152,3 +152,100 @@ suite "Token Bucket": reqTime += 1.minutes let (maxCompensationBudget, _, _) = bucket.getAvailableCapacity(reqTime) check maxCompensationBudget == 1250 # 1000 + 250 max compensation + + test "TokenBucket custom budget parameter": + let now = Moment.now() + + # Test with default budget (-1, should use budgetCap) + var bucket1 = TokenBucket.new(1000, 1.seconds, budget = -1) + let (budget1, budgetCap1, _) = bucket1.getAvailableCapacity(now) + check budget1 == 1000 + check budgetCap1 == 1000 + + # Test with custom budget less than capacity + var bucket2 = TokenBucket.new(1000, 1.seconds, budget = 500) + let (budget2, budgetCap2, _) = bucket2.getAvailableCapacity(now) + check budget2 == 500 + check budgetCap2 == 1000 + + # Test with budget equal to capacity + var bucket3 = TokenBucket.new(1000, 1.seconds, budget = 1000) + let (budget3, budgetCap3, _) = bucket3.getAvailableCapacity(now) + check budget3 == 1000 + check budgetCap3 == 1000 + + # Test with zero budget + var bucket4 = TokenBucket.new(1000, 1.seconds, budget = 0) + let (budget4, budgetCap4, _) = bucket4.getAvailableCapacity(now) + check budget4 == 0 + check budgetCap4 == 1000 + + # Test consumption with custom budget + check bucket2.tryConsume(300, now) == true + let (budget2After, budgetCap2After, _) = bucket2.getAvailableCapacity(now) + check budget2After == 200 + check budgetCap2After == 1000 + check bucket2.tryConsume(300, now) == false # Should fail, only 200 remaining + + test "TokenBucket custom lastTimeFull parameter": + let now = Moment.now() + let pastTime = now - 5.seconds + + # Test with past lastTimeFull (bucket should be ready for refill) + var bucket1 = TokenBucket.new(1000, 1.seconds, budget = 0, lastTimeFull = pastTime) + # Since 5 seconds have passed and period is 1 second, bucket should refill + check bucket1.tryConsume(1000, now) == true + + # Test with current time as lastTimeFull + var bucket2 = TokenBucket.new(1000, 1.seconds, budget = 500, lastTimeFull = now) + check bucket2.getAvailableCapacity(now) == (500, 1000, now) + + # Test strict mode with past lastTimeFull + var bucket3 = TokenBucket.new( + 1000, 1.seconds, ReplenishMode.Strict, budget = 0, lastTimeFull = pastTime + ) + check bucket3.tryConsume(1000, now) == true # Should refill to full capacity + + # Test compensating mode with past lastTimeFull and partial usage + var bucket4 = TokenBucket.new( + 1000, 2.seconds, ReplenishMode.Compensating, budget = 300, lastTimeFull = pastTime + ) + # 5 seconds passed, period is 2 seconds, so 2.5 periods elapsed + # Usage average = 300/1000/2.5 = 0.12 (12% usage) + # Compensation = (1.0 - 0.12) * 1000 = 880, capped at 250 + let (availableBudget, _, _) = bucket4.getAvailableCapacity(now) + check availableBudget == 1250 # 1000 + 250 max compensation + + test "TokenBucket parameter combinations": + let now = Moment.now() + let pastTime = now - 3.seconds + + # Test strict mode with custom budget and past time + var strictBucket = TokenBucket.new( + budgetCap = 500, + fillDuration = 1.seconds, + mode = ReplenishMode.Strict, + budget = 100, + lastTimeFull = pastTime, + ) + check strictBucket.getAvailableCapacity(now) == (500, 500, pastTime) + # Should show full capacity since period elapsed + + # Test compensating mode with custom budget and past time + var compBucket = TokenBucket.new( + budgetCap = 1000, + fillDuration = 2.seconds, + mode = ReplenishMode.Compensating, + budget = 200, + lastTimeFull = pastTime, + ) + # 3 seconds passed, period is 2 seconds, so 1.5 periods elapsed + # Usage average = 200/1000/1.5 = 0.133 (13.3% usage) + # Compensation = (1.0 - 0.133) * 1000 = 867, capped at 250 + let (compensatedBudget, _, _) = compBucket.getAvailableCapacity(now) + check compensatedBudget == 1250 # 1000 + 250 max compensation + + # Test edge case: zero budget with immediate consumption + var zeroBucket = TokenBucket.new(100, 1.seconds, budget = 0, lastTimeFull = now) + check zeroBucket.tryConsume(1, now) == false # Should fail immediately + check zeroBucket.tryConsume(1, now + 1.seconds) == true # Should succeed after period