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.