From edefa13aebe0dc7b07f6e7e1526a81d3c179c2dd Mon Sep 17 00:00:00 2001 From: NagyZoltanPeter <113987313+NagyZoltanPeter@users.noreply.github.com> Date: Sun, 24 Aug 2025 15:05:15 +0200 Subject: [PATCH 1/8] Fix original TokenBucket, consolidate the behaviors of chronos/TokenBucket and waku/TokenBucket (compensating) with no interface change --- chronos/ratelimit.nim | 49 ++++++++++++++++++++++++++----------- tests/testratelimit.nim | 54 ++++++++++++++++++++++++++++++++++++----- 2 files changed, 83 insertions(+), 20 deletions(-) diff --git a/chronos/ratelimit.nim b/chronos/ratelimit.nim index ad66c067e..d35214393 100644 --- a/chronos/ratelimit.nim +++ b/chronos/ratelimit.nim @@ -33,26 +33,47 @@ proc update(bucket: TokenBucket, currentTime: Moment) = bucket.budget = min(bucket.budgetCap, bucket.budget) return - if currentTime < bucket.lastUpdate: + if currentTime <= bucket.lastUpdate: return - let - timeDelta = currentTime - bucket.lastUpdate - fillPercent = timeDelta.milliseconds.float / bucket.fillDuration.milliseconds.float - replenished = - int(bucket.budgetCap.float * fillPercent) - deltaFromReplenished = - int(bucket.fillDuration.milliseconds.float * - replenished.float / bucket.budgetCap.float) + let timeDelta = currentTime - bucket.lastUpdate + let cap = bucket.budgetCap + let periodNs = bucket.fillDuration.nanoseconds.int64 + let deltaNs = timeDelta.nanoseconds.int64 - bucket.lastUpdate += milliseconds(deltaFromReplenished) - bucket.budget = min(bucket.budgetCap, bucket.budget + replenished) + # How many whole tokens could be produced by the elapsed time. + let possibleTokens = int((deltaNs * cap.int64) div periodNs) + if possibleTokens <= 0: + return + + let budgetLeft = cap - bucket.budget + if budgetLeft <= 0: + # Bucket already full the entire elapsed time: burn the elapsed time + # so we do not accumulate implicit credit. + bucket.lastUpdate = currentTime + bucket.budget = cap # do not let over budgeting + return + + let toAdd = min(possibleTokens, budgetLeft) + + # Advance lastUpdate only by the fraction of time actually “spent” to mint toAdd tokens. + # (toAdd / cap) * period = time used + let usedNs = (periodNs * toAdd.int64) div cap.int64 + bucket.budget += toAdd + if toAdd == budgetLeft and possibleTokens > budgetLeft: + # We hit the cap; discard leftover elapsed time to prevent multi-call burst inflation + bucket.lastUpdate = currentTime + else: + bucket.lastUpdate += nanoseconds(usedNs) proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool = ## If `tokens` are available, consume them, - ## Otherwhise, return false. + ## Otherwise, return false. if bucket.budget >= tokens: + # If bucket was full, burn elapsed time to avoid immediate refill + flake. + if bucket.budget == bucket.budgetCap: + bucket.lastUpdate = now bucket.budget -= tokens return true @@ -60,9 +81,9 @@ proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool = if bucket.budget >= tokens: bucket.budget -= tokens - true + return true else: - false + return false proc worker(bucket: TokenBucket) {.async.} = while bucket.pendingRequests.len > 0: diff --git a/tests/testratelimit.nim b/tests/testratelimit.nim index d28492874..872ac9541 100644 --- a/tests/testratelimit.nim +++ b/tests/testratelimit.nim @@ -117,14 +117,56 @@ suite "Token Bucket": check bucket.tryConsume(1, fakeNow) == true test "Short replenish": - skip() + # skip() # TODO (cheatfate): This test was disabled, because it continuosly fails in # Github Actions Windows x64 CI when using Nim 1.6.14 version. # Unable to reproduce failure locally. - # var bucket = TokenBucket.new(15000, 1.milliseconds) - # let start = Moment.now() - # check bucket.tryConsume(15000, start) - # check bucket.tryConsume(1, start) == false + var bucket = TokenBucket.new(15000, 1.milliseconds) + let start = Moment.now() + check bucket.tryConsume(15000, start) + check bucket.tryConsume(1, start) == false - # check bucket.tryConsume(15000, start + 1.milliseconds) == true + check bucket.tryConsume(15000, start + 1.milliseconds) == true + + # Edge-case: ensure only one refill can occur for the same timestamp even if + # multiple tryConsume calls are made that would otherwise appear to have large + # elapsed time credit. This prevents multi-call burst inflation at a single time. + test "No double refill at same timestamp": + var bucket = TokenBucket.new(10, 100.milliseconds) + let t0 = Moment.now() + # Consume from full so lastUpdate is stamped at t0 + check bucket.tryConsume(5, t0) == true # budget now 5 + # Long idle period (simulate large elapsed time) + let idle = t0 + 5.seconds + # First large request triggers an update + refill limited by space (5) + check bucket.tryConsume(6, idle) == true # budget after = 4 (5 minted -> 10 then -6) + # Second request at the SAME timestamp cannot refill again + check bucket.tryConsume(5, idle) == false + # Prove only 4 remain: consuming 4 succeeds, then 1 more fails at same timestamp + check bucket.tryConsume(4, idle) == true + check bucket.tryConsume(1, idle) == false + + # Edge-case fairness: partial usage should only mint up to available space, not + # more than cap, and leftover elapsed time is burned once cap is reached. + test "Refill limited by available space": + var bucket = TokenBucket.new(10, 100.milliseconds) + let t0 = Moment.now() + # Spend a portion (from full) -> lastUpdate = t0, budget 4 + check bucket.tryConsume(6, t0) == true + # Mid-period small consume without triggering update (still before refill point) + let mid = t0 + 50.milliseconds + check bucket.tryConsume(1, mid) == true # budget 3 + # At the 100ms boundary request more than remaining budget to force update + let boundary = t0 + 100.milliseconds + # Space is 7; even though 100ms elapsed corresponds to 10 possible tokens, + # only 7 are minted and leftover elapsed time credit is discarded. + check bucket.tryConsume(6, boundary) == true # leaves 4 + # A second consume at identical boundary timestamp cannot mint more than residual + check bucket.tryConsume(5, boundary) == false + # After another 40ms, at most floor(40/100 * 10)=4 tokens accrue; request 4 succeeds + let late = boundary + 40.milliseconds + check bucket.tryConsume(4, late) == true # should deplete + # A subsequent call at the same timestamp may mint remaining fractional time credit (fair catch-up) + # so a small consume can still succeed. + check bucket.tryConsume(1, late) == true From d63a1bc68c69dbc711308e03912ff48c34f0d044 Mon Sep 17 00:00:00 2001 From: NagyZoltanPeter <113987313+NagyZoltanPeter@users.noreply.github.com> Date: Sun, 24 Aug 2025 15:49:51 +0200 Subject: [PATCH 2/8] TokenBucket extended with waku's strict mode replenish, that does not allow refill just after period boundary elapsed. --- chronos/ratelimit.nim | 33 ++++++++++++++++++++++++++++++--- tests/testratelimit.nim | 10 ++++++++++ 2 files changed, 40 insertions(+), 3 deletions(-) diff --git a/chronos/ratelimit.nim b/chronos/ratelimit.nim index d35214393..4cacddeeb 100644 --- a/chronos/ratelimit.nim +++ b/chronos/ratelimit.nim @@ -14,6 +14,10 @@ import timer export timer type + ReplenishMode* = enum + Strict + Ballanced + BucketWaiter = object future: Future[void] value: int @@ -27,8 +31,23 @@ type workFuture: Future[void] pendingRequests: seq[BucketWaiter] manuallyReplenished: AsyncEvent + replenishMode: ReplenishMode -proc update(bucket: TokenBucket, currentTime: Moment) = +func fullPeriodElapsed(bucket: TokenBucket, currentTime: Moment): bool = + return currentTime - bucket.lastUpdate >= bucket.fillDuration + +proc updateStrict(bucket: TokenBucket, currentTime: Moment) = + if bucket.fillDuration == default(Duration): + bucket.budget = min(bucket.budgetCap, bucket.budget) + return + + if not fullPeriodElapsed(bucket, currentTime): + return + + bucket.budget = bucket.budgetCap + bucket.lastUpdate = currentTime + +proc updateBallanced(bucket: TokenBucket, currentTime: Moment) = if bucket.fillDuration == default(Duration): bucket.budget = min(bucket.budgetCap, bucket.budget) return @@ -66,6 +85,12 @@ proc update(bucket: TokenBucket, currentTime: Moment) = else: bucket.lastUpdate += nanoseconds(usedNs) +proc update(bucket: TokenBucket, currentTime: Moment) = + if bucket.replenishMode == ReplenishMode.Strict: + updateStrict(bucket, currentTime) + else: + updateBallanced(bucket, currentTime) + proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool = ## If `tokens` are available, consume them, ## Otherwise, return false. @@ -151,7 +176,8 @@ proc replenish*(bucket: TokenBucket, tokens: int, now = Moment.now()) = proc new*( T: type[TokenBucket], budgetCap: int, - fillDuration: Duration = 1.seconds): T = + fillDuration: Duration = 1.seconds, + replenishMode: ReplenishMode = ReplenishMode.Ballanced): T = ## Create a TokenBucket T( @@ -159,5 +185,6 @@ proc new*( budgetCap: budgetCap, fillDuration: fillDuration, lastUpdate: Moment.now(), - manuallyReplenished: newAsyncEvent() + manuallyReplenished: newAsyncEvent(), + replenishMode: replenishMode ) diff --git a/tests/testratelimit.nim b/tests/testratelimit.nim index 872ac9541..d14ff3316 100644 --- a/tests/testratelimit.nim +++ b/tests/testratelimit.nim @@ -170,3 +170,13 @@ suite "Token Bucket": # A subsequent call at the same timestamp may mint remaining fractional time credit (fair catch-up) # so a small consume can still succeed. check bucket.tryConsume(1, late) == true + + test "Strict replenish mode does not refill before period elapsed": + var bucket = TokenBucket.new(10, 100.milliseconds, ReplenishMode.Strict) + let t0 = Moment.now() + # Spend a portion (from full) -> lastUpdate = t0, budget 4 + check bucket.tryConsume(9, t0) == true # leaves 1 + let mid = t0 + 50.milliseconds + check bucket.tryConsume(2, mid) == false # budget 1 + let boundary = t0 + 100.milliseconds + check bucket.tryConsume(2, boundary) == true # leaves 8 From eec4f0303c0fa0a1c990e16dbc4de7e601664c5e Mon Sep 17 00:00:00 2001 From: NagyZoltanPeter <113987313+NagyZoltanPeter@users.noreply.github.com> Date: Mon, 25 Aug 2025 00:10:02 +0200 Subject: [PATCH 3/8] Adjust for new chat-sdk needs, added setState and getAvailableCapacity with small refactoring --- chronos/ratelimit.nim | 60 ++++++++++++++++++++++++++--------------- tests/testratelimit.nim | 19 +++++++++++++ 2 files changed, 57 insertions(+), 22 deletions(-) diff --git a/chronos/ratelimit.nim b/chronos/ratelimit.nim index 4cacddeeb..e436ecdcc 100644 --- a/chronos/ratelimit.nim +++ b/chronos/ratelimit.nim @@ -36,24 +36,21 @@ type func fullPeriodElapsed(bucket: TokenBucket, currentTime: Moment): bool = return currentTime - bucket.lastUpdate >= bucket.fillDuration -proc updateStrict(bucket: TokenBucket, currentTime: Moment) = +proc calcUpdateStrict(bucket: TokenBucket, currentTime: Moment): tuple[budget: int, lastUpdate: Moment] = if bucket.fillDuration == default(Duration): - bucket.budget = min(bucket.budgetCap, bucket.budget) - return + return (min(bucket.budgetCap, bucket.budget), bucket.lastUpdate) if not fullPeriodElapsed(bucket, currentTime): - return + return (bucket.budget, bucket.lastUpdate) - bucket.budget = bucket.budgetCap - bucket.lastUpdate = currentTime + return (bucket.budgetCap, currentTime) -proc updateBallanced(bucket: TokenBucket, currentTime: Moment) = +proc calcUpdateBallanced(bucket: TokenBucket, currentTime: Moment): tuple[budget: int, lastUpdate: Moment] = if bucket.fillDuration == default(Duration): - bucket.budget = min(bucket.budgetCap, bucket.budget) - return + return (min(bucket.budgetCap, bucket.budget), bucket.lastUpdate) if currentTime <= bucket.lastUpdate: - return + return (bucket.budget, bucket.lastUpdate) let timeDelta = currentTime - bucket.lastUpdate let cap = bucket.budgetCap @@ -63,33 +60,37 @@ proc updateBallanced(bucket: TokenBucket, currentTime: Moment) = # How many whole tokens could be produced by the elapsed time. let possibleTokens = int((deltaNs * cap.int64) div periodNs) if possibleTokens <= 0: - return + return (bucket.budget, bucket.lastUpdate) let budgetLeft = cap - bucket.budget if budgetLeft <= 0: # Bucket already full the entire elapsed time: burn the elapsed time - # so we do not accumulate implicit credit. - bucket.lastUpdate = currentTime - bucket.budget = cap # do not let over budgeting - return + # so we do not accumulate implicit credit and do not allow over budgeting + return (cap, currentTime) let toAdd = min(possibleTokens, budgetLeft) # Advance lastUpdate only by the fraction of time actually “spent” to mint toAdd tokens. # (toAdd / cap) * period = time used let usedNs = (periodNs * toAdd.int64) div cap.int64 - bucket.budget += toAdd + let newbudget = bucket.budget + toAdd + var newLastUpdate = bucket.lastUpdate + nanoseconds(usedNs) if toAdd == budgetLeft and possibleTokens > budgetLeft: # We hit the cap; discard leftover elapsed time to prevent multi-call burst inflation - bucket.lastUpdate = currentTime - else: - bucket.lastUpdate += nanoseconds(usedNs) + newLastUpdate = currentTime -proc update(bucket: TokenBucket, currentTime: Moment) = + return (newbudget, newLastUpdate) + +proc calcUpdate(bucket: TokenBucket, currentTime: Moment): tuple[budget: int, lastUpdate: Moment] = if bucket.replenishMode == ReplenishMode.Strict: - updateStrict(bucket, currentTime) + return bucket.calcUpdateStrict(currentTime) else: - updateBallanced(bucket, currentTime) + return bucket.calcUpdateBallanced(currentTime) + +proc update(bucket: TokenBucket, currentTime: Moment) = + let (newBudget, newLastUpdate) = bucket.calcUpdate(currentTime) + bucket.budget = newBudget + bucket.lastUpdate = newLastUpdate proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool = ## If `tokens` are available, consume them, @@ -173,6 +174,12 @@ proc replenish*(bucket: TokenBucket, tokens: int, now = Moment.now()) = bucket.update(now) bucket.manuallyReplenished.fire() +proc getAvailableCapacity*( + bucket: TokenBucket, currentTime: Moment = Moment.now() +): tuple[budget: int, budgetCap: int, lastUpdate: Moment] = + let (assumedBudget, assumedLastUpdate) = bucket.calcUpdate(currentTime) + return (assumedBudget, bucket.budgetCap, assumedLastUpdate) + proc new*( T: type[TokenBucket], budgetCap: int, @@ -188,3 +195,12 @@ proc new*( manuallyReplenished: newAsyncEvent(), replenishMode: replenishMode ) + +proc setState*(bucket: TokenBucket, budget: int, lastUpdate: Moment) = + bucket.budget = budget + bucket.lastUpdate = lastUpdate + +func `$`*(b: TokenBucket): string {.inline.} = + if isNil(b): + return "nil" + return $b.budgetCap & "/" & $b.fillDuration diff --git a/tests/testratelimit.nim b/tests/testratelimit.nim index d14ff3316..f8585259e 100644 --- a/tests/testratelimit.nim +++ b/tests/testratelimit.nim @@ -176,7 +176,26 @@ suite "Token Bucket": let t0 = Moment.now() # Spend a portion (from full) -> lastUpdate = t0, budget 4 check bucket.tryConsume(9, t0) == true # leaves 1 + + var cap = bucket.getAvailableCapacity(t0) + check cap.budget == 1 + check cap.lastUpdate == t0 + check cap.budgetCap == 10 + let mid = t0 + 50.milliseconds + + cap = bucket.getAvailableCapacity(mid) + check cap.budget == 1 + check cap.lastUpdate == t0 + check cap.budgetCap == 10 + check bucket.tryConsume(2, mid) == false # budget 1 + let boundary = t0 + 100.milliseconds + + cap = bucket.getAvailableCapacity(boundary) + check cap.budget == 10 + check cap.lastUpdate == boundary + check cap.budgetCap == 10 + check bucket.tryConsume(2, boundary) == true # leaves 8 From a7f008fbb04c445e71288e34267ea84b7d013185 Mon Sep 17 00:00:00 2001 From: NagyZoltanPeter <113987313+NagyZoltanPeter@users.noreply.github.com> Date: Mon, 25 Aug 2025 00:45:09 +0200 Subject: [PATCH 4/8] Better comments --- chronos/ratelimit.nim | 2 +- tests/testratelimit.nim | 11 +++-------- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/chronos/ratelimit.nim b/chronos/ratelimit.nim index e436ecdcc..599c9efff 100644 --- a/chronos/ratelimit.nim +++ b/chronos/ratelimit.nim @@ -97,7 +97,7 @@ proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool = ## Otherwise, return false. if bucket.budget >= tokens: - # If bucket was full, burn elapsed time to avoid immediate refill + flake. + # If bucket is full, consider this point as period start, drop silent periods before if bucket.budget == bucket.budgetCap: bucket.lastUpdate = now bucket.budget -= tokens diff --git a/tests/testratelimit.nim b/tests/testratelimit.nim index f8585259e..db6462444 100644 --- a/tests/testratelimit.nim +++ b/tests/testratelimit.nim @@ -117,11 +117,6 @@ suite "Token Bucket": check bucket.tryConsume(1, fakeNow) == true test "Short replenish": - # skip() - # TODO (cheatfate): This test was disabled, because it continuosly fails in - # Github Actions Windows x64 CI when using Nim 1.6.14 version. - # Unable to reproduce failure locally. - var bucket = TokenBucket.new(15000, 1.milliseconds) let start = Moment.now() check bucket.tryConsume(15000, start) @@ -174,7 +169,7 @@ suite "Token Bucket": test "Strict replenish mode does not refill before period elapsed": var bucket = TokenBucket.new(10, 100.milliseconds, ReplenishMode.Strict) let t0 = Moment.now() - # Spend a portion (from full) -> lastUpdate = t0, budget 4 + # Spend a portion (from full) -> lastUpdate = t0, budget 10 check bucket.tryConsume(9, t0) == true # leaves 1 var cap = bucket.getAvailableCapacity(t0) @@ -189,7 +184,7 @@ suite "Token Bucket": check cap.lastUpdate == t0 check cap.budgetCap == 10 - check bucket.tryConsume(2, mid) == false # budget 1 + check bucket.tryConsume(2, mid) == false # no update before period boundary passed, budget 1 let boundary = t0 + 100.milliseconds @@ -198,4 +193,4 @@ suite "Token Bucket": check cap.lastUpdate == boundary check cap.budgetCap == 10 - check bucket.tryConsume(2, boundary) == true # leaves 8 + check bucket.tryConsume(2, boundary) == true # ok, we passed the period boundary now, leaves 8 From 1c85269576c47389f6e74a671963937a2679f2ea Mon Sep 17 00:00:00 2001 From: NagyZoltanPeter <113987313+NagyZoltanPeter@users.noreply.github.com> Date: Mon, 25 Aug 2025 14:24:45 +0200 Subject: [PATCH 5/8] rename for better self explain code, added more explanatory comments upon code review finding --- chronos/ratelimit.nim | 55 +++++++++++++++++++++++------------------ tests/testratelimit.nim | 6 ++--- 2 files changed, 34 insertions(+), 27 deletions(-) diff --git a/chronos/ratelimit.nim b/chronos/ratelimit.nim index 599c9efff..5c9c10f72 100644 --- a/chronos/ratelimit.nim +++ b/chronos/ratelimit.nim @@ -15,8 +15,12 @@ export timer type ReplenishMode* = enum + # Strict mode allows replenish tokens only after the fill duration has elapsed. Strict - Ballanced + + # For better utilization of available tokens and tolerate burst periods. + # Balanced mode allows minting tokens based on elapsed time in between consuming of tokens up to budget capacity. + Balanced BucketWaiter = object future: Future[void] @@ -25,7 +29,7 @@ type TokenBucket* = ref object budget: int - budgetCap: int + budgetCapacity: int lastUpdate: Moment fillDuration: Duration workFuture: Future[void] @@ -38,45 +42,48 @@ func fullPeriodElapsed(bucket: TokenBucket, currentTime: Moment): bool = proc calcUpdateStrict(bucket: TokenBucket, currentTime: Moment): tuple[budget: int, lastUpdate: Moment] = if bucket.fillDuration == default(Duration): - return (min(bucket.budgetCap, bucket.budget), bucket.lastUpdate) + # with zero fillDuration we only allow manual replenish till capacity + return (min(bucket.budgetCapacity, bucket.budget), bucket.lastUpdate) if not fullPeriodElapsed(bucket, currentTime): return (bucket.budget, bucket.lastUpdate) - return (bucket.budgetCap, currentTime) + return (bucket.budgetCapacity, currentTime) -proc calcUpdateBallanced(bucket: TokenBucket, currentTime: Moment): tuple[budget: int, lastUpdate: Moment] = +proc calcUpdateBalanced(bucket: TokenBucket, currentTime: Moment): tuple[budget: int, lastUpdate: Moment] = if bucket.fillDuration == default(Duration): - return (min(bucket.budgetCap, bucket.budget), bucket.lastUpdate) + # with zero fillDuration we only allow manual replenish till capacity + return (min(bucket.budgetCapacity, bucket.budget), bucket.lastUpdate) if currentTime <= bucket.lastUpdate: + # don't allow backward timing return (bucket.budget, bucket.lastUpdate) let timeDelta = currentTime - bucket.lastUpdate - let cap = bucket.budgetCap + let capacity = bucket.budgetCapacity let periodNs = bucket.fillDuration.nanoseconds.int64 let deltaNs = timeDelta.nanoseconds.int64 # How many whole tokens could be produced by the elapsed time. - let possibleTokens = int((deltaNs * cap.int64) div periodNs) + let possibleTokens = int((deltaNs * capacity.int64) div periodNs) if possibleTokens <= 0: return (bucket.budget, bucket.lastUpdate) - let budgetLeft = cap - bucket.budget + let budgetLeft = capacity - bucket.budget if budgetLeft <= 0: # Bucket already full the entire elapsed time: burn the elapsed time # so we do not accumulate implicit credit and do not allow over budgeting - return (cap, currentTime) + return (capacity, currentTime) let toAdd = min(possibleTokens, budgetLeft) # Advance lastUpdate only by the fraction of time actually “spent” to mint toAdd tokens. - # (toAdd / cap) * period = time used - let usedNs = (periodNs * toAdd.int64) div cap.int64 + # (toAdd / capacity) * period = time used + let usedNs = (periodNs * toAdd.int64) div capacity.int64 let newbudget = bucket.budget + toAdd var newLastUpdate = bucket.lastUpdate + nanoseconds(usedNs) if toAdd == budgetLeft and possibleTokens > budgetLeft: - # We hit the cap; discard leftover elapsed time to prevent multi-call burst inflation + # We hit the capacity; discard leftover elapsed time to prevent multi-call burst inflation newLastUpdate = currentTime return (newbudget, newLastUpdate) @@ -85,7 +92,7 @@ proc calcUpdate(bucket: TokenBucket, currentTime: Moment): tuple[budget: int, la if bucket.replenishMode == ReplenishMode.Strict: return bucket.calcUpdateStrict(currentTime) else: - return bucket.calcUpdateBallanced(currentTime) + return bucket.calcUpdateBalanced(currentTime) proc update(bucket: TokenBucket, currentTime: Moment) = let (newBudget, newLastUpdate) = bucket.calcUpdate(currentTime) @@ -98,7 +105,7 @@ proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool = if bucket.budget >= tokens: # If bucket is full, consider this point as period start, drop silent periods before - if bucket.budget == bucket.budgetCap: + if bucket.budget == bucket.budgetCapacity: bucket.lastUpdate = now bucket.budget -= tokens return true @@ -127,8 +134,8 @@ proc worker(bucket: TokenBucket) {.async.} = let eventWaiter = bucket.manuallyReplenished.wait() if bucket.fillDuration.milliseconds > 0: let - nextCycleValue = float(min(waiter.value, bucket.budgetCap)) - budgetRatio = nextCycleValue.float / bucket.budgetCap.float + nextCycleValue = float(min(waiter.value, bucket.budgetCapacity)) + budgetRatio = nextCycleValue.float / bucket.budgetCapacity.float timeToTarget = int(budgetRatio * bucket.fillDuration.milliseconds.float) + 1 #TODO this will create a timer for each blocked bucket, #which may cause performance issue when creating many @@ -176,20 +183,20 @@ proc replenish*(bucket: TokenBucket, tokens: int, now = Moment.now()) = proc getAvailableCapacity*( bucket: TokenBucket, currentTime: Moment = Moment.now() -): tuple[budget: int, budgetCap: int, lastUpdate: Moment] = +): tuple[budget: int, budgetCapacity: int, lastUpdate: Moment] = let (assumedBudget, assumedLastUpdate) = bucket.calcUpdate(currentTime) - return (assumedBudget, bucket.budgetCap, assumedLastUpdate) + return (assumedBudget, bucket.budgetCapacity, assumedLastUpdate) proc new*( T: type[TokenBucket], - budgetCap: int, + budgetCapacity: int, fillDuration: Duration = 1.seconds, - replenishMode: ReplenishMode = ReplenishMode.Ballanced): T = + replenishMode: ReplenishMode = ReplenishMode.Balanced): T = ## Create a TokenBucket T( - budget: budgetCap, - budgetCap: budgetCap, + budget: budgetCapacity, + budgetCapacity: budgetCapacity, fillDuration: fillDuration, lastUpdate: Moment.now(), manuallyReplenished: newAsyncEvent(), @@ -203,4 +210,4 @@ proc setState*(bucket: TokenBucket, budget: int, lastUpdate: Moment) = func `$`*(b: TokenBucket): string {.inline.} = if isNil(b): return "nil" - return $b.budgetCap & "/" & $b.fillDuration + return $b.budgetCapacity & "/" & $b.fillDuration diff --git a/tests/testratelimit.nim b/tests/testratelimit.nim index db6462444..cf7ab5c37 100644 --- a/tests/testratelimit.nim +++ b/tests/testratelimit.nim @@ -175,14 +175,14 @@ suite "Token Bucket": var cap = bucket.getAvailableCapacity(t0) check cap.budget == 1 check cap.lastUpdate == t0 - check cap.budgetCap == 10 + check cap.budgetCapacity == 10 let mid = t0 + 50.milliseconds cap = bucket.getAvailableCapacity(mid) check cap.budget == 1 check cap.lastUpdate == t0 - check cap.budgetCap == 10 + check cap.budgetCapacity == 10 check bucket.tryConsume(2, mid) == false # no update before period boundary passed, budget 1 @@ -191,6 +191,6 @@ suite "Token Bucket": cap = bucket.getAvailableCapacity(boundary) check cap.budget == 10 check cap.lastUpdate == boundary - check cap.budgetCap == 10 + check cap.budgetCapacity == 10 check bucket.tryConsume(2, boundary) == true # ok, we passed the period boundary now, leaves 8 From 451ecd69740bb45917d27625b626a9424a7d09b6 Mon Sep 17 00:00:00 2001 From: NagyZoltanPeter <113987313+NagyZoltanPeter@users.noreply.github.com> Date: Fri, 19 Sep 2025 12:55:10 +0200 Subject: [PATCH 6/8] Document TokenBucket with detailed samples of Balanced mode replenishment algorithm, extended TokenBucket unit test --- docs/src/ratelimit.md | 171 ++++++++++++++++++++++++++++++ tests/testratelimit.nim | 228 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 399 insertions(+) create mode 100644 docs/src/ratelimit.md diff --git a/docs/src/ratelimit.md b/docs/src/ratelimit.md new file mode 100644 index 000000000..d09abe4b5 --- /dev/null +++ b/docs/src/ratelimit.md @@ -0,0 +1,171 @@ +## TokenBucket — Usage Modes (Overview) + +TokenBucket provides several usage modes and patterns depending on how you want to rate-limit: + +- Balanced mode (default): + - Mints tokens proportionally to elapsed time at a constant rate (`capacity / fillDuration`), adding only whole tokens. + - When the bucket is full for an interval, the elapsed time is burned (no “credit banking”). + - If an update would overfill, budget is clamped to capacity and leftover elapsed time is discarded; `lastUpdate` is set to the current time. + - Nanosecond-level accounting for precise behavior. + +- Strict mode: + - Replenishes only after a full `fillDuration` has elapsed (step-like refill behavior). + - Before the period boundary, budget does not increase; after the boundary, budget jumps to capacity. + - Use when you need hard period boundaries rather than proportional accrual. + +- Manual-only replenish (fillDuration = 0): + - Disables automatic minting; tokens can only be added via `replenish(tokens)`. + - Replenish is capped at capacity and wakes pending consumers. + +- Synchronous consumption: `tryConsume(tokens, now)` + - Attempts to consume immediately; returns `true` on success, `false` otherwise. + - If consuming from full, `lastUpdate` is set to `now` (prevents idle-at-full credit banking in Balanced mode). + +- Asynchronous consumption: `consume(tokens, now) -> Future[void]` + - Returns a future that completes when tokens become available (or can be cancelled). + - Internally, the waiter is woken around the time enough tokens are expected to accrue, or earlier if `replenish()` is called. + +- Capacity and timing introspection: `getAvailableCapacity(now)` + - Computes the budget as of `now` without mutating bucket state. + +- Manual replenishment: `replenish(tokens, now)` + - Adds tokens (capped to capacity), updates timing, and wakes waiters. + +The sections below illustrate Balanced semantics with concrete timelines and compare them with the older algorithm for context. + +# TokenBucket Balanced Mode — Scenario 1 Timeline + +Assumptions: +- Capacity `C = 10` +- `fillDuration = 1s` (per-token time: 100ms) +- Start: `t = 0ms`, `budget = 10`, `lastUpdate = 0ms` + +Legend: +- Minted tokens: tokens added by Balanced update at that step (TA) +- Budget after mint: budget after minting, before the consume at that row +- Budget after consume: budget left after processing the request at that row +- LU set?: whether `lastUpdate` changes at that step (reason) + +Only request events are listed below (no passive availability checks): + +| Time | Elapsed from LU | Budget (in) | Request tokens | Minted tokens (TA) | Budget after mint | Budget after consume | LU set? | +|---------|------------------|-------------|----------------|--------------------|-------------------|----------------------|-----------------------------------| +| 0 ms | n/a | 10 | 7 | 0 | 10 | 3 | yes (consume/full → 0 ms) | +| 200 ms | 200 ms | 3 | 5 | 2 | 5 | 0 | yes (update → 200 ms) | +| 650 ms | 450 ms | 0 | 3 | 4 | 4 | 1 | yes (update → 600 ms) | +| 1200 ms | 600 ms | 1 | 6 | 6 | 7 | 1 | yes (update → 1200 ms) | +| 1800 ms | 600 ms | 1 | 5 | 6 | 7 | 2 | yes (update → 1800 ms) | +| 2100 ms | 300 ms | 2 | 10 | 3 | 5 | 5 (insufficient) | yes (update → 2100 ms) | +| 2600 ms | 500 ms | 5 | 10 | 5 (to cap) | 10 (hit cap) | 0 | yes (update hit cap → 2600 ms); yes (consume/full → 2600 ms) | + +Notes: +- When an update would overfill the bucket, it is clamped to capacity and `lastUpdate` is set to the current time; leftover elapsed time is discarded. +- Consuming from a full bucket sets `lastUpdate` to the consume time (prevents idle-at-full credit banking). + +### Consumption Summary (0–3s window) + +Per `fillDuration` period (1s each): + +| Period | Requests within period | Tokens consumed | +|----------------|-------------------------------------------------|-----------------| +| 0–1000 ms | 0ms:7, 200ms:5, 650ms:3 | 15 | +| 1000–2000 ms | 1200ms:6, 1800ms:5 | 11 | +| 2000–3000 ms | 2100ms:10 (insufficient), 2600ms:10 (consumed) | 10 | + +Total consumed over 3 seconds: 15 + 11 + 10 = 36 tokens. + +## Old Algorithm (Pre-unification) — Scenario 1 Timeline + +Reference: old `TokenBucket` from master (`chronos/ratelimit.nim`) before the unification. + +Key behavioral differences vs current Balanced mode: +- No LU reset when consuming from a full bucket (LU stays unchanged on consume). +- No explicit burn of leftover elapsed time when capacity is reached; fractional leftover time is retained via LU based on minted tokens. + +Assumptions and inputs are identical to the table above: +- Capacity `C = 10`, `fillDuration = 1s` (100ms per token) +- Request-only events at: 0ms (7), 200ms (5), 650ms (3), 1200ms (6), 1800ms (5), 2100ms (10), 2600ms (10) +- Start: `t = 0ms`, `budget = 10`, `lastUpdate = 0ms` + +| Time | Elapsed from LU | Budget (in) | Request tokens | Minted tokens (TA) | Budget after mint | Budget after consume | LU set? | +|---------|------------------|-------------|----------------|--------------------|-------------------|----------------------|------------------------------| +| 0 ms | n/a | 10 | 7 | 0 | 10 | 3 | no (consume does not set LU) | +| 200 ms | 200 ms | 3 | 5 | 2 | 5 | 0 | yes (update → 200 ms) | +| 650 ms | 450 ms | 0 | 3 | 4 | 4 | 1 | yes (update → 600 ms) | +| 1200 ms | 600 ms | 1 | 6 | 6 | 7 | 1 | yes (update → 1200 ms) | +| 1800 ms | 600 ms | 1 | 5 | 6 | 7 | 2 | yes (update → 1800 ms) | +| 2100 ms | 300 ms | 2 | 10 | 3 | 5 | 5 (insufficient) | yes (update → 2100 ms) | +| 2600 ms | 500 ms | 5 | 10 | 5 (to cap) | 10 (hit cap) | 0 | yes (update → 2600 ms) | + +### Consumption Summary (0–3s window, old algorithm) + +| Period | Requests within period | Tokens consumed | +|----------------|-------------------------------------------------|-----------------| +| 0–1000 ms | 0ms:7, 200ms:5, 650ms:3 | 15 | +| 1000–2000 ms | 1200ms:6, 1800ms:5 | 11 | +| 2000–3000 ms | 2100ms:10 (insufficient), 2600ms:10 (consumed) | 10 | + +Total consumed over 3 seconds (old): 36 tokens. + +## Diff Notes — Balanced vs Old + +- Consume from full: + - Balanced: sets `lastUpdate` to the consume time before subtracting. + - Old: does not modify `lastUpdate` on consume. + +- When the bucket is full during an elapsed interval: + - Balanced: burns all the elapsed time by setting `lastUpdate = currentTime` (no hidden time credit). + - Old: advances `lastUpdate` only by the time used to mint whole tokens, leaving fractional leftover time to carry to the next update. + +- Hitting capacity mid-update (overfill path): + - Balanced: clamps to capacity and sets `lastUpdate = currentTime`, discarding leftover elapsed time. + - Old: clamps to capacity but sets `lastUpdate = lastUpdate + usedTime` (time to mint the whole tokens), retaining the leftover fractional part of the elapsed interval. + +- Time resolution: + - Balanced: nanosecond-based accounting. + - Old: millisecond-based accounting. Sub-ms differences and rounding may diverge. + +Practical impact: +- Balanced avoids multi-call burst inflation at a single timestamp and prevents banking implicit credit during idle-at-full periods. +- In this scenario with millisecond-aligned times, total consumption over 3 seconds is the same (36 tokens), but LU handling differs and becomes visible with sub-token leftover time, long idle while full, or overfill updates. + +### Comparison Examples + +| Case | Inputs | Balanced outcome | Old outcome | Key difference | +|-------------------------------|------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------| +| Sub-token leftover time | C=10, T=1s, budget=4, LU=1.000s; check at t=1.075s (Δ=75ms) | PT=0 minted; budget stays 4; LU unchanged | PT=0 minted; budget stays 4; LU unchanged | No mint below one-token time; both unchanged (ns vs ms resolution doesn’t change the outcome) | +| Long idle while full | C=5, T=1s, budget=5 (full), LU=0; next update at t=2.5s (Δ=2500ms) | Budget remains 5; LU set to 2.5s (burn entire idle interval; no leftover fractional time) | Replenished=floor(5×2.5)=12 (clamped); used=12×200ms=2400ms; budget=5; LU=2.4s (keeps 100ms leftover) | Balanced burns idle-at-full time; Old retains fractional leftover (100ms) | +| Overfull update (hit capacity)| C=10, T=1s, budget=8 (space=2), LU=0; update at t=300ms (Δ=300ms) | PT=floor(3)=3; TA=min(3,2)=2; budget→10; LU=300ms (hit-cap path discards leftover 100ms) | PT=3; TA=2; budget→10; used=2×100ms=200ms; LU=200ms (leftover 100ms retained) | Balanced discards leftover on cap-hit; Old keeps leftover time | + +## High-rate single-token requests (Balanced) + +Settings: +- Capacity `C = 10`, `fillDuration = 10ms` (per-token time: 1ms) +- Window to observe: `0–40ms` (4 full periods) +- Requests are 1 token each; batches occur at specific timestamps. + +We show how the bucket rejects attempts that exceed the available budget at each instant, ensuring no more than `capacity + minted` tokens are usable in any time frame. Over `0–40ms`, at most `10 (initial capacity) + 4 × 10 (mint) = 50` tokens can be consumed. + +Request batches and outcomes: + +| Time | Elapsed from LU | Budget before | Minted (PT→TA) | Budget after mint | Requests (×1) | Accepted | Rejected | Budget after consume | LU after | +|--------|------------------|---------------|-----------------|-------------------|---------------|----------|----------|----------------------|---------| +| 0 ms | n/a | 10 | 0 | 10 | 12 | 10 | 2 | 0 | 0 ms | +| 5 ms | 5 ms | 0 | 5 → 5 | 5 | 7 | 5 | 2 | 0 | 5 ms | +| 10 ms | 5 ms | 0 | 5 → 5 | 5 | 15 | 5 | 10 | 0 | 10 ms | +| 12 ms | 2 ms | 0 | 2 → 2 | 2 | 3 | 2 | 1 | 0 | 12 ms | +| 20 ms | 8 ms | 0 | 8 → 8 | 8 | 25 | 8 | 17 | 0 | 20 ms | +| 30 ms | 10 ms | 0 | 10 → 10 | 10 | 9 | 9 | 0 | 1 | 30 ms | +| 31 ms | 1 ms | 1 | 1 → 1 | 2 | 3 | 2 | 1 | 0 | 31 ms | +| 40 ms | 9 ms | 0 | 9 → 9 | 9 | 20 | 9 | 11 | 0 | 40 ms | + +Totals over 0–40ms: +- Attempted: 12 + 7 + 15 + 3 + 25 + 9 + 3 + 20 = 94 requests +- Accepted: 10 + 5 + 5 + 2 + 8 + 9 + 2 + 9 = 50 tokens (matches `10 + 4×10`) +- Rejected: 94 − 50 = 44 requests + +Why the rejections happen (preventing overuse): +- At any given instant, you can only consume up to the tokens currently in the bucket. +- Between instants, tokens mint continuously at `capacity / fillDuration = 1 token/ms`; the table shows how many become available just before each batch. +- When a batch demands more than available, the excess is rejected (or would be queued with `consume()`), enforcing the rate limit. +- Over any observation window, the maximum consumable tokens = initial available (up to capacity) + tokens minted during that window; here, that cap is `10 + (40ms × 1/ms) = 50`. \ No newline at end of file diff --git a/tests/testratelimit.nim b/tests/testratelimit.nim index cf7ab5c37..42fa314bc 100644 --- a/tests/testratelimit.nim +++ b/tests/testratelimit.nim @@ -194,3 +194,231 @@ suite "Token Bucket": check cap.budgetCapacity == 10 check bucket.tryConsume(2, boundary) == true # ok, we passed the period boundary now, leaves 8 + + test "Balanced high-rate single-token 10/10ms over 40ms": + # Capacity 10, fillDuration 10ms (1 token/ms). Only 1-token requests are made + # at specific timestamps within 0–40ms (4 full periods). We verify that no + # more than 50 tokens can be consumed in total and that per-batch accept/reject + # counts and lastUpdate values match expectations. + var bucket = TokenBucket.new(10, 10.milliseconds) + let t0 = Moment.now() + let t5 = t0 + 5.milliseconds + let t10 = t0 + 10.milliseconds + let t12 = t0 + 12.milliseconds + let t20 = t0 + 20.milliseconds + let t30 = t0 + 30.milliseconds + let t31 = t0 + 31.milliseconds + let t40 = t0 + 40.milliseconds + + proc attempt(count: int, now: Moment): tuple[accepted, rejected: int] = + var acc = 0 + var rej = 0 + for _ in 0.. accept 10, reject 2; budget ends 0; LU=0ms (consume-from-full) + var r = attempt(12, t0) + check r.accepted == 10 + check r.rejected == 2 + var cap = bucket.getAvailableCapacity(t0) + check cap.budget == 0 + check cap.lastUpdate == t0 + + # 5ms: 7 attempts -> mint 5 then accept 5, reject 2; budget 0; LU=5ms + r = attempt(7, t5) + check r.accepted == 5 + check r.rejected == 2 + cap = bucket.getAvailableCapacity(t5) + check cap.budget == 0 + check cap.lastUpdate == t5 + + # 10ms: 15 attempts -> mint 5 then accept 5, reject 10; budget 0; LU=10ms + r = attempt(15, t10) + check r.accepted == 5 + check r.rejected == 10 + cap = bucket.getAvailableCapacity(t10) + check cap.budget == 0 + check cap.lastUpdate == t10 + + # 12ms: 3 attempts -> mint 2 then accept 2, reject 1; budget 0; LU=12ms + r = attempt(3, t12) + check r.accepted == 2 + check r.rejected == 1 + cap = bucket.getAvailableCapacity(t12) + check cap.budget == 0 + check cap.lastUpdate == t12 + + # 20ms: 25 attempts -> mint 8 then accept 8, reject 17; budget 0; LU=20ms + r = attempt(25, t20) + check r.accepted == 8 + check r.rejected == 17 + cap = bucket.getAvailableCapacity(t20) + check cap.budget == 0 + check cap.lastUpdate == t20 + + # 30ms: 9 attempts -> mint 10 then accept 9, budget ends 1; LU=30ms + r = attempt(9, t30) + check r.accepted == 9 + check r.rejected == 0 + cap = bucket.getAvailableCapacity(t30) + check cap.budget == 1 + check cap.lastUpdate == t30 + + # 31ms: 3 attempts -> mint 1 then accept 2, reject 1; budget 0; LU=31ms + r = attempt(3, t31) + check r.accepted == 2 + check r.rejected == 1 + cap = bucket.getAvailableCapacity(t31) + check cap.budget == 0 + check cap.lastUpdate == t31 + + # 40ms: 20 attempts -> mint 9 then accept 9, reject 11; budget 0; LU=40ms + r = attempt(20, t40) + check r.accepted == 9 + check r.rejected == 11 + cap = bucket.getAvailableCapacity(t40) + check cap.budget == 0 + check cap.lastUpdate == t40 + + # Totals across 0–40ms window + let totalAccepted = 10 + 5 + 5 + 2 + 8 + 9 + 2 + 9 + let totalRejected = 2 + 2 + 10 + 1 + 17 + 0 + 1 + 11 + check totalAccepted == 50 + check totalRejected == 44 + + test "Balanced high-rate single-token 10/10ms over 40ms (advancing time)": + # Variant of the high-rate test where each tryConsume occurs at a timestamp that + # advances by ~1ms when possible, simulating a more realistic stream of requests. + # We still demand more than can be provided to ensure rejections, and we verify + # that across 0–40ms only 50 tokens can be accepted. + var bucket = TokenBucket.new(10, 10.milliseconds) + let t0 = Moment.now() + + var accepted = 0 + var rejected = 0 + # Perform 94 attempts; for the first 41 attempts we increase the timestamp by 1ms + # per attempt (up to t0+40ms). After that, we keep attempting at t0+40ms to + # simulate concurrent bursts at the end of the window. + for i in 0..<94: + let ts = t0 + min(i, 40).milliseconds + if bucket.tryConsume(1, ts): + inc accepted + else: + inc rejected + + # At most 10 (initial) + 40 (minted over 40ms) can be accepted + check accepted == 50 + check rejected == 44 + + let t40 = t0 + 40.milliseconds + let cap = bucket.getAvailableCapacity(t40) + # All available tokens within the window should have been consumed by our attempts + check cap.budget == 0 + + # Balanced-mode scenario reproductions and timeline validation + test "Balanced Scenario 1 timeline": + # Capacity 10, fillDuration 1s, per-token time 100ms + var bucket = TokenBucket.new(10, 1.seconds) + + let t0 = Moment.now() + let t200 = t0 + 200.milliseconds + let t600 = t0 + 600.milliseconds + let t650 = t0 + 650.milliseconds + let t1000 = t0 + 1000.milliseconds + let t1200 = t0 + 1200.milliseconds + let t1800 = t0 + 1800.milliseconds + let t2100 = t0 + 2100.milliseconds + let t2600 = t0 + 2600.milliseconds + let t3000 = t0 + 3000.milliseconds + + # 1) t=0ms: consume 7 from full -> LU set to t0, budget 3 + check bucket.tryConsume(7, t0) == true + var cap = bucket.getAvailableCapacity(t0) + check cap.budget == 3 + check cap.lastUpdate == t0 + + # 2) t=200ms: request 5 -> mint 2, then consume -> budget 0, LU=200ms + check bucket.tryConsume(5, t200) == true + cap = bucket.getAvailableCapacity(t200) + check cap.budget == 0 + check cap.lastUpdate == t200 + + # 3) t=650ms: request 3 -> mint 4 (to t600), then consume -> budget 1, LU=600ms + check bucket.tryConsume(3, t650) == true + cap = bucket.getAvailableCapacity(t600) + check cap.budget == 1 + check cap.lastUpdate == t600 + + # 4) t=1000ms: availability check only (does not mutate); expected 4 minted -> budget 5, LU=1000ms + cap = bucket.getAvailableCapacity(t1000) + check cap.budget == 5 + check cap.lastUpdate == t1000 + + # 5) t=1200ms: request 6 -> net minted since LU to here totals 6 -> budget ends 1, LU=1200ms + check bucket.tryConsume(6, t1200) == true + cap = bucket.getAvailableCapacity(t1200) + check cap.budget == 1 + check cap.lastUpdate == t1200 + + # 6) t=1800ms: request 5 -> mint 6 then consume -> budget 2, LU=1800ms + check bucket.tryConsume(5, t1800) == true + cap = bucket.getAvailableCapacity(t1800) + check cap.budget == 2 + check cap.lastUpdate == t1800 + + # 7) t=2100ms: request 10 -> mint 3 to reach 5, still insufficient -> false, LU=2100ms + check bucket.tryConsume(10, t2100) == false + cap = bucket.getAvailableCapacity(t2100) + check cap.budget == 5 + check cap.lastUpdate == t2100 + + # 8) t=2600ms: enough time for 5 more -> reach full and then consume 10 -> budget 0, LU=2600ms + check bucket.tryConsume(10, t2600) == true + cap = bucket.getAvailableCapacity(t2600) + check cap.budget == 0 + check cap.lastUpdate == t2600 + + # 9) t=3000ms: availability check -> mint 4 -> budget 4, LU=3000ms + cap = bucket.getAvailableCapacity(t3000) + check cap.budget == 4 + check cap.lastUpdate == t3000 + + test "Balanced: idle while full burns time": + # Capacity 5, fillDuration 1s, per-token time 200ms + var bucket = TokenBucket.new(5, 1.seconds) + let t0 = Moment.now() + let t2_5 = t0 + 2500.milliseconds + # Consume 1 after long idle at full -> LU set to now + check bucket.tryConsume(1, t2_5) == true + var cap = bucket.getAvailableCapacity(t2_5) + check cap.budget == 4 + check cap.lastUpdate == t2_5 + # 100ms later is below per-token time -> no mint + let t2_6 = t0 + 2600.milliseconds + cap = bucket.getAvailableCapacity(t2_6) + check cap.budget == 4 + check cap.lastUpdate == t2_5 + + test "Balanced: large jump clamps to capacity and LU=now when capped": + # Capacity 8, fillDuration 4s, per-token time 0.5s + var bucket = TokenBucket.new(8, 4.seconds) + let t0 = Moment.now() + # Spend 6 from full so LU = t0, budget = 2 + check bucket.tryConsume(6, t0) == true + + let t5 = t0 + 5.seconds + # Availability check: should reach cap and set LU to t5 (hit-cap path burns leftover time) + var cap2 = bucket.getAvailableCapacity(t5) + check cap2.budget == 8 + check cap2.lastUpdate == t5 + + # Consume 3 slightly later; update will also clamp and set LU to now, then consume from full + let t5_2 = t0 + 5200.milliseconds + check bucket.tryConsume(3, t5_2) == true + cap2 = bucket.getAvailableCapacity(t5_2) + check cap2.budget == 5 + check cap2.lastUpdate == t5_2 From 7f65929111ae8df548a25d9c7e3f51bb3efa2680 Mon Sep 17 00:00:00 2001 From: NagyZoltanPeter <113987313+NagyZoltanPeter@users.noreply.github.com> Date: Tue, 21 Oct 2025 02:52:48 +0200 Subject: [PATCH 7/8] Address review comments, renaming the replenish modes, removed old/new algo comparison from documentation, fix Discrete mode update calculation to properly calculate correct last update time by period distance calculation. --- chronos/ratelimit.nim | 67 +++++++++++++++++----------------- docs/src/ratelimit.md | 80 +++++------------------------------------ tests/testratelimit.nim | 22 ++++++------ 3 files changed, 55 insertions(+), 114 deletions(-) diff --git a/chronos/ratelimit.nim b/chronos/ratelimit.nim index 5c9c10f72..d9985fe7b 100644 --- a/chronos/ratelimit.nim +++ b/chronos/ratelimit.nim @@ -14,13 +14,11 @@ import timer export timer type - ReplenishMode* = enum - # Strict mode allows replenish tokens only after the fill duration has elapsed. - Strict - - # For better utilization of available tokens and tolerate burst periods. - # Balanced mode allows minting tokens based on elapsed time in between consuming of tokens up to budget capacity. - Balanced + ReplenishMode* = enum + Continuous + # Tokens are continuously replenished at a rate of `capacity / fillDuration`, up to the configured capacity + Discrete + # Up to `capacity` tokens are replenished once every `fillDuration`, in discrete steps, such that at the beginning of every `fillDuration` period, there are `capacity` tokens available BucketWaiter = object future: Future[void] @@ -29,7 +27,7 @@ type TokenBucket* = ref object budget: int - budgetCapacity: int + capacity: int lastUpdate: Moment fillDuration: Duration workFuture: Future[void] @@ -37,30 +35,35 @@ type manuallyReplenished: AsyncEvent replenishMode: ReplenishMode -func fullPeriodElapsed(bucket: TokenBucket, currentTime: Moment): bool = - return currentTime - bucket.lastUpdate >= bucket.fillDuration +func periodDistance(bucket: TokenBucket, currentTime: Moment): float = + if currentTime <= bucket.lastUpdate: + return 0.0 + + return nanoseconds(currentTime - bucket.lastUpdate).float / + nanoseconds(bucket.fillDuration).float -proc calcUpdateStrict(bucket: TokenBucket, currentTime: Moment): tuple[budget: int, lastUpdate: Moment] = +proc calcUpdateDiscrete(bucket: TokenBucket, currentTime: Moment): tuple[budget: int, lastUpdate: Moment] = if bucket.fillDuration == default(Duration): # with zero fillDuration we only allow manual replenish till capacity - return (min(bucket.budgetCapacity, bucket.budget), bucket.lastUpdate) + return (min(bucket.capacity, bucket.budget), bucket.lastUpdate) - if not fullPeriodElapsed(bucket, currentTime): + let distance = periodDistance(bucket, currentTime) + if distance < 1.0: return (bucket.budget, bucket.lastUpdate) - return (bucket.budgetCapacity, currentTime) + return (bucket.capacity, bucket.lastUpdate + (bucket.fillDuration * int(distance))) -proc calcUpdateBalanced(bucket: TokenBucket, currentTime: Moment): tuple[budget: int, lastUpdate: Moment] = +proc calcUpdateContinuous(bucket: TokenBucket, currentTime: Moment): tuple[budget: int, lastUpdate: Moment] = if bucket.fillDuration == default(Duration): # with zero fillDuration we only allow manual replenish till capacity - return (min(bucket.budgetCapacity, bucket.budget), bucket.lastUpdate) + return (min(bucket.capacity, bucket.budget), bucket.lastUpdate) if currentTime <= bucket.lastUpdate: # don't allow backward timing return (bucket.budget, bucket.lastUpdate) let timeDelta = currentTime - bucket.lastUpdate - let capacity = bucket.budgetCapacity + let capacity = bucket.capacity let periodNs = bucket.fillDuration.nanoseconds.int64 let deltaNs = timeDelta.nanoseconds.int64 @@ -89,10 +92,10 @@ proc calcUpdateBalanced(bucket: TokenBucket, currentTime: Moment): tuple[budget: return (newbudget, newLastUpdate) proc calcUpdate(bucket: TokenBucket, currentTime: Moment): tuple[budget: int, lastUpdate: Moment] = - if bucket.replenishMode == ReplenishMode.Strict: - return bucket.calcUpdateStrict(currentTime) + if bucket.replenishMode == ReplenishMode.Discrete: + return bucket.calcUpdateDiscrete(currentTime) else: - return bucket.calcUpdateBalanced(currentTime) + return bucket.calcUpdateContinuous(currentTime) proc update(bucket: TokenBucket, currentTime: Moment) = let (newBudget, newLastUpdate) = bucket.calcUpdate(currentTime) @@ -105,7 +108,7 @@ proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool = if bucket.budget >= tokens: # If bucket is full, consider this point as period start, drop silent periods before - if bucket.budget == bucket.budgetCapacity: + if bucket.budget == bucket.capacity: bucket.lastUpdate = now bucket.budget -= tokens return true @@ -114,9 +117,9 @@ proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool = if bucket.budget >= tokens: bucket.budget -= tokens - return true + true else: - return false + false proc worker(bucket: TokenBucket) {.async.} = while bucket.pendingRequests.len > 0: @@ -134,8 +137,8 @@ proc worker(bucket: TokenBucket) {.async.} = let eventWaiter = bucket.manuallyReplenished.wait() if bucket.fillDuration.milliseconds > 0: let - nextCycleValue = float(min(waiter.value, bucket.budgetCapacity)) - budgetRatio = nextCycleValue.float / bucket.budgetCapacity.float + nextCycleValue = float(min(waiter.value, bucket.capacity)) + budgetRatio = nextCycleValue.float / bucket.capacity.float timeToTarget = int(budgetRatio * bucket.fillDuration.milliseconds.float) + 1 #TODO this will create a timer for each blocked bucket, #which may cause performance issue when creating many @@ -183,20 +186,20 @@ proc replenish*(bucket: TokenBucket, tokens: int, now = Moment.now()) = proc getAvailableCapacity*( bucket: TokenBucket, currentTime: Moment = Moment.now() -): tuple[budget: int, budgetCapacity: int, lastUpdate: Moment] = +): tuple[budget: int, capacity: int, lastUpdate: Moment] = let (assumedBudget, assumedLastUpdate) = bucket.calcUpdate(currentTime) - return (assumedBudget, bucket.budgetCapacity, assumedLastUpdate) + return (assumedBudget, bucket.capacity, assumedLastUpdate) proc new*( T: type[TokenBucket], - budgetCapacity: int, + capacity: int, fillDuration: Duration = 1.seconds, - replenishMode: ReplenishMode = ReplenishMode.Balanced): T = + replenishMode: ReplenishMode = ReplenishMode.Continuous): T = ## Create a TokenBucket T( - budget: budgetCapacity, - budgetCapacity: budgetCapacity, + budget: capacity, + capacity: capacity, fillDuration: fillDuration, lastUpdate: Moment.now(), manuallyReplenished: newAsyncEvent(), @@ -210,4 +213,4 @@ proc setState*(bucket: TokenBucket, budget: int, lastUpdate: Moment) = func `$`*(b: TokenBucket): string {.inline.} = if isNil(b): return "nil" - return $b.budgetCapacity & "/" & $b.fillDuration + return $b.capacity & "/" & $b.fillDuration diff --git a/docs/src/ratelimit.md b/docs/src/ratelimit.md index d09abe4b5..796312300 100644 --- a/docs/src/ratelimit.md +++ b/docs/src/ratelimit.md @@ -1,14 +1,14 @@ -## TokenBucket — Usage Modes (Overview) +# TokenBucket — Usage Modes (Overview) TokenBucket provides several usage modes and patterns depending on how you want to rate-limit: -- Balanced mode (default): +- Continuous mode (default): - Mints tokens proportionally to elapsed time at a constant rate (`capacity / fillDuration`), adding only whole tokens. - When the bucket is full for an interval, the elapsed time is burned (no “credit banking”). - If an update would overfill, budget is clamped to capacity and leftover elapsed time is discarded; `lastUpdate` is set to the current time. - Nanosecond-level accounting for precise behavior. -- Strict mode: +- Discrete mode: - Replenishes only after a full `fillDuration` has elapsed (step-like refill behavior). - Before the period boundary, budget does not increase; after the boundary, budget jumps to capacity. - Use when you need hard period boundaries rather than proportional accrual. @@ -19,7 +19,7 @@ TokenBucket provides several usage modes and patterns depending on how you want - Synchronous consumption: `tryConsume(tokens, now)` - Attempts to consume immediately; returns `true` on success, `false` otherwise. - - If consuming from full, `lastUpdate` is set to `now` (prevents idle-at-full credit banking in Balanced mode). + - If consuming from full, `lastUpdate` is set to `now` (prevents idle-at-full credit banking in Continuous mode). - Asynchronous consumption: `consume(tokens, now) -> Future[void]` - Returns a future that completes when tokens become available (or can be cancelled). @@ -31,9 +31,10 @@ TokenBucket provides several usage modes and patterns depending on how you want - Manual replenishment: `replenish(tokens, now)` - Adds tokens (capped to capacity), updates timing, and wakes waiters. -The sections below illustrate Balanced semantics with concrete timelines and compare them with the older algorithm for context. +The sections below illustrate Continuous semantics with concrete timelines and compare them with the older algorithm for context. -# TokenBucket Balanced Mode — Scenario 1 Timeline +# Example Scenarios for Continuous Mode +## TokenBucket Continuous Mode — Scenario 1 Timeline Assumptions: - Capacity `C = 10` @@ -41,7 +42,7 @@ Assumptions: - Start: `t = 0ms`, `budget = 10`, `lastUpdate = 0ms` Legend: -- Minted tokens: tokens added by Balanced update at that step (TA) +- Minted tokens: tokens added by Continuous update at that step (TA) - Budget after mint: budget after minting, before the consume at that row - Budget after consume: budget left after processing the request at that row - LU set?: whether `lastUpdate` changes at that step (reason) @@ -74,70 +75,7 @@ Per `fillDuration` period (1s each): Total consumed over 3 seconds: 15 + 11 + 10 = 36 tokens. -## Old Algorithm (Pre-unification) — Scenario 1 Timeline - -Reference: old `TokenBucket` from master (`chronos/ratelimit.nim`) before the unification. - -Key behavioral differences vs current Balanced mode: -- No LU reset when consuming from a full bucket (LU stays unchanged on consume). -- No explicit burn of leftover elapsed time when capacity is reached; fractional leftover time is retained via LU based on minted tokens. - -Assumptions and inputs are identical to the table above: -- Capacity `C = 10`, `fillDuration = 1s` (100ms per token) -- Request-only events at: 0ms (7), 200ms (5), 650ms (3), 1200ms (6), 1800ms (5), 2100ms (10), 2600ms (10) -- Start: `t = 0ms`, `budget = 10`, `lastUpdate = 0ms` - -| Time | Elapsed from LU | Budget (in) | Request tokens | Minted tokens (TA) | Budget after mint | Budget after consume | LU set? | -|---------|------------------|-------------|----------------|--------------------|-------------------|----------------------|------------------------------| -| 0 ms | n/a | 10 | 7 | 0 | 10 | 3 | no (consume does not set LU) | -| 200 ms | 200 ms | 3 | 5 | 2 | 5 | 0 | yes (update → 200 ms) | -| 650 ms | 450 ms | 0 | 3 | 4 | 4 | 1 | yes (update → 600 ms) | -| 1200 ms | 600 ms | 1 | 6 | 6 | 7 | 1 | yes (update → 1200 ms) | -| 1800 ms | 600 ms | 1 | 5 | 6 | 7 | 2 | yes (update → 1800 ms) | -| 2100 ms | 300 ms | 2 | 10 | 3 | 5 | 5 (insufficient) | yes (update → 2100 ms) | -| 2600 ms | 500 ms | 5 | 10 | 5 (to cap) | 10 (hit cap) | 0 | yes (update → 2600 ms) | - -### Consumption Summary (0–3s window, old algorithm) - -| Period | Requests within period | Tokens consumed | -|----------------|-------------------------------------------------|-----------------| -| 0–1000 ms | 0ms:7, 200ms:5, 650ms:3 | 15 | -| 1000–2000 ms | 1200ms:6, 1800ms:5 | 11 | -| 2000–3000 ms | 2100ms:10 (insufficient), 2600ms:10 (consumed) | 10 | - -Total consumed over 3 seconds (old): 36 tokens. - -## Diff Notes — Balanced vs Old - -- Consume from full: - - Balanced: sets `lastUpdate` to the consume time before subtracting. - - Old: does not modify `lastUpdate` on consume. - -- When the bucket is full during an elapsed interval: - - Balanced: burns all the elapsed time by setting `lastUpdate = currentTime` (no hidden time credit). - - Old: advances `lastUpdate` only by the time used to mint whole tokens, leaving fractional leftover time to carry to the next update. - -- Hitting capacity mid-update (overfill path): - - Balanced: clamps to capacity and sets `lastUpdate = currentTime`, discarding leftover elapsed time. - - Old: clamps to capacity but sets `lastUpdate = lastUpdate + usedTime` (time to mint the whole tokens), retaining the leftover fractional part of the elapsed interval. - -- Time resolution: - - Balanced: nanosecond-based accounting. - - Old: millisecond-based accounting. Sub-ms differences and rounding may diverge. - -Practical impact: -- Balanced avoids multi-call burst inflation at a single timestamp and prevents banking implicit credit during idle-at-full periods. -- In this scenario with millisecond-aligned times, total consumption over 3 seconds is the same (36 tokens), but LU handling differs and becomes visible with sub-token leftover time, long idle while full, or overfill updates. - -### Comparison Examples - -| Case | Inputs | Balanced outcome | Old outcome | Key difference | -|-------------------------------|------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------| -| Sub-token leftover time | C=10, T=1s, budget=4, LU=1.000s; check at t=1.075s (Δ=75ms) | PT=0 minted; budget stays 4; LU unchanged | PT=0 minted; budget stays 4; LU unchanged | No mint below one-token time; both unchanged (ns vs ms resolution doesn’t change the outcome) | -| Long idle while full | C=5, T=1s, budget=5 (full), LU=0; next update at t=2.5s (Δ=2500ms) | Budget remains 5; LU set to 2.5s (burn entire idle interval; no leftover fractional time) | Replenished=floor(5×2.5)=12 (clamped); used=12×200ms=2400ms; budget=5; LU=2.4s (keeps 100ms leftover) | Balanced burns idle-at-full time; Old retains fractional leftover (100ms) | -| Overfull update (hit capacity)| C=10, T=1s, budget=8 (space=2), LU=0; update at t=300ms (Δ=300ms) | PT=floor(3)=3; TA=min(3,2)=2; budget→10; LU=300ms (hit-cap path discards leftover 100ms) | PT=3; TA=2; budget→10; used=2×100ms=200ms; LU=200ms (leftover 100ms retained) | Balanced discards leftover on cap-hit; Old keeps leftover time | - -## High-rate single-token requests (Balanced) +## High-rate single-token requests (Continuous) Settings: - Capacity `C = 10`, `fillDuration = 10ms` (per-token time: 1ms) diff --git a/tests/testratelimit.nim b/tests/testratelimit.nim index 42fa314bc..f8a34f69f 100644 --- a/tests/testratelimit.nim +++ b/tests/testratelimit.nim @@ -166,8 +166,8 @@ suite "Token Bucket": # so a small consume can still succeed. check bucket.tryConsume(1, late) == true - test "Strict replenish mode does not refill before period elapsed": - var bucket = TokenBucket.new(10, 100.milliseconds, ReplenishMode.Strict) + test "Discrete replenish mode does not refill before period elapsed": + var bucket = TokenBucket.new(10, 100.milliseconds, ReplenishMode.Discrete) let t0 = Moment.now() # Spend a portion (from full) -> lastUpdate = t0, budget 10 check bucket.tryConsume(9, t0) == true # leaves 1 @@ -175,14 +175,14 @@ suite "Token Bucket": var cap = bucket.getAvailableCapacity(t0) check cap.budget == 1 check cap.lastUpdate == t0 - check cap.budgetCapacity == 10 + check cap.capacity == 10 let mid = t0 + 50.milliseconds cap = bucket.getAvailableCapacity(mid) check cap.budget == 1 check cap.lastUpdate == t0 - check cap.budgetCapacity == 10 + check cap.capacity == 10 check bucket.tryConsume(2, mid) == false # no update before period boundary passed, budget 1 @@ -191,11 +191,11 @@ suite "Token Bucket": cap = bucket.getAvailableCapacity(boundary) check cap.budget == 10 check cap.lastUpdate == boundary - check cap.budgetCapacity == 10 + check cap.capacity == 10 check bucket.tryConsume(2, boundary) == true # ok, we passed the period boundary now, leaves 8 - test "Balanced high-rate single-token 10/10ms over 40ms": + test "Continuous high-rate single-token 10/10ms over 40ms": # Capacity 10, fillDuration 10ms (1 token/ms). Only 1-token requests are made # at specific timestamps within 0–40ms (4 full periods). We verify that no # more than 50 tokens can be consumed in total and that per-batch accept/reject @@ -290,7 +290,7 @@ suite "Token Bucket": check totalAccepted == 50 check totalRejected == 44 - test "Balanced high-rate single-token 10/10ms over 40ms (advancing time)": + test "Continuous high-rate single-token 10/10ms over 40ms (advancing time)": # Variant of the high-rate test where each tryConsume occurs at a timestamp that # advances by ~1ms when possible, simulating a more realistic stream of requests. # We still demand more than can be provided to ensure rejections, and we verify @@ -319,8 +319,8 @@ suite "Token Bucket": # All available tokens within the window should have been consumed by our attempts check cap.budget == 0 - # Balanced-mode scenario reproductions and timeline validation - test "Balanced Scenario 1 timeline": + # Continuous-mode scenario reproductions and timeline validation + test "Continuous Scenario 1 timeline": # Capacity 10, fillDuration 1s, per-token time 100ms var bucket = TokenBucket.new(10, 1.seconds) @@ -387,7 +387,7 @@ suite "Token Bucket": check cap.budget == 4 check cap.lastUpdate == t3000 - test "Balanced: idle while full burns time": + test "Continuous: idle while full burns time": # Capacity 5, fillDuration 1s, per-token time 200ms var bucket = TokenBucket.new(5, 1.seconds) let t0 = Moment.now() @@ -403,7 +403,7 @@ suite "Token Bucket": check cap.budget == 4 check cap.lastUpdate == t2_5 - test "Balanced: large jump clamps to capacity and LU=now when capped": + test "Continuous: large jump clamps to capacity and LU=now when capped": # Capacity 8, fillDuration 4s, per-token time 0.5s var bucket = TokenBucket.new(8, 4.seconds) let t0 = Moment.now() From b9cf3d1e115d83eb2fcc278094feb6e396874f7d Mon Sep 17 00:00:00 2001 From: NagyZoltanPeter <113987313+NagyZoltanPeter@users.noreply.github.com> Date: Fri, 24 Oct 2025 16:51:17 +0200 Subject: [PATCH 8/8] protect from devide by zero, code style fix --- chronos/ratelimit.nim | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/chronos/ratelimit.nim b/chronos/ratelimit.nim index d9985fe7b..508e6a262 100644 --- a/chronos/ratelimit.nim +++ b/chronos/ratelimit.nim @@ -36,11 +36,10 @@ type replenishMode: ReplenishMode func periodDistance(bucket: TokenBucket, currentTime: Moment): float = - if currentTime <= bucket.lastUpdate: + if currentTime <= bucket.lastUpdate or bucket.fillDuration == default(Duration): return 0.0 - return nanoseconds(currentTime - bucket.lastUpdate).float / - nanoseconds(bucket.fillDuration).float + nanoseconds(currentTime - bucket.lastUpdate).float / nanoseconds(bucket.fillDuration).float proc calcUpdateDiscrete(bucket: TokenBucket, currentTime: Moment): tuple[budget: int, lastUpdate: Moment] = if bucket.fillDuration == default(Duration): @@ -51,7 +50,7 @@ proc calcUpdateDiscrete(bucket: TokenBucket, currentTime: Moment): tuple[budget: if distance < 1.0: return (bucket.budget, bucket.lastUpdate) - return (bucket.capacity, bucket.lastUpdate + (bucket.fillDuration * int(distance))) + (bucket.capacity, bucket.lastUpdate + (bucket.fillDuration * int(distance))) proc calcUpdateContinuous(bucket: TokenBucket, currentTime: Moment): tuple[budget: int, lastUpdate: Moment] = if bucket.fillDuration == default(Duration): @@ -89,7 +88,7 @@ proc calcUpdateContinuous(bucket: TokenBucket, currentTime: Moment): tuple[budge # We hit the capacity; discard leftover elapsed time to prevent multi-call burst inflation newLastUpdate = currentTime - return (newbudget, newLastUpdate) + (newbudget, newLastUpdate) proc calcUpdate(bucket: TokenBucket, currentTime: Moment): tuple[budget: int, lastUpdate: Moment] = if bucket.replenishMode == ReplenishMode.Discrete: @@ -176,7 +175,7 @@ proc consume*(bucket: TokenBucket, tokens: int, now = Moment.now()): Future[void if isNil(bucket.workFuture) or bucket.workFuture.finished(): bucket.workFuture = worker(bucket) - return retFuture + retFuture proc replenish*(bucket: TokenBucket, tokens: int, now = Moment.now()) = ## Add `tokens` to the budget (capped to the bucket capacity) @@ -188,7 +187,7 @@ proc getAvailableCapacity*( bucket: TokenBucket, currentTime: Moment = Moment.now() ): tuple[budget: int, capacity: int, lastUpdate: Moment] = let (assumedBudget, assumedLastUpdate) = bucket.calcUpdate(currentTime) - return (assumedBudget, bucket.capacity, assumedLastUpdate) + (assumedBudget, bucket.capacity, assumedLastUpdate) proc new*( T: type[TokenBucket], @@ -213,4 +212,4 @@ proc setState*(bucket: TokenBucket, budget: int, lastUpdate: Moment) = func `$`*(b: TokenBucket): string {.inline.} = if isNil(b): return "nil" - return $b.capacity & "/" & $b.fillDuration + $b.capacity & "/" & $b.fillDuration