From fe54bb26d3addb8435b61adc0c4fc945fd7c1b2d Mon Sep 17 00:00:00 2001 From: pablo Date: Thu, 17 Jul 2025 11:31:31 +0300 Subject: [PATCH] initial --- .gitignore | 24 ++++ Makefile | 56 ++++++++ README.md | 34 +++++ tests/test_token_bucket.nim | 251 ++++++++++++++++++++++++++++++++++ token_bucket.nimble | 9 ++ token_bucket/token_bucket.nim | 204 +++++++++++++++++++++++++++ 6 files changed, 578 insertions(+) create mode 100644 .gitignore create mode 100644 Makefile create mode 100644 README.md create mode 100644 tests/test_token_bucket.nim create mode 100644 token_bucket.nimble create mode 100644 token_bucket/token_bucket.nim diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e3f957f --- /dev/null +++ b/.gitignore @@ -0,0 +1,24 @@ +# IDE +.vscode/ + +# Nim +nimcache/ +*.exe +*.dll +*.so +*.dylib + +# Build artifacts +*.o +*.a + +# OS +.DS_Store + +# Compiled files +token_bucket/* +!*.nim +tests/* +!*.nim +nimble.develop +nimble.paths \ No newline at end of file diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..ad765ea --- /dev/null +++ b/Makefile @@ -0,0 +1,56 @@ +# Token Bucket Makefile +# Build system for the Nim token bucket rate limiting library + +.PHONY: all build test clean install deps help + +# Default target +all: deps build test + +# Help target +help: + @echo "Token Bucket Build System" + @echo "=========================" + @echo "Available targets:" + @echo " all - Install deps, build, and test" + @echo " build - Build the library" + @echo " test - Run tests" + @echo " deps - Install dependencies" + @echo " clean - Clean build artifacts" + @echo " help - Show this help message" + +# Install dependencies +deps: + @echo "Installing dependencies..." + nimble install -d + +# Build the library +build: + @echo "Building token bucket library..." + nim c --mm:refc token_bucket/token_bucket.nim + +# Run tests +test: + @echo "Running tests..." + nim c -r tests/test_token_bucket.nim + +# Clean build artifacts +clean: + @echo "Cleaning build artifacts..." + rm -rf token_bucket/nimcache tests/nimcache + rm -f token_bucket/token_bucket tests/test_token_bucket + find . -name "*.pdb" -delete 2>/dev/null || true + +# Install the package +install: + @echo "Installing package..." + nimble install + +# Development build with debug info +debug: + @echo "Building with debug info..." + nim c --debuginfo --linedir:on token_bucket/token_bucket.nim + +# Run tests with coverage (if available) +test-coverage: + @echo "Running tests with coverage..." + nim c -r --passC:"-fprofile-arcs -ftest-coverage" --passL:"-lgcov" tests/test_token_bucket.nim \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..5ff41ab --- /dev/null +++ b/README.md @@ -0,0 +1,34 @@ +# Token Bucket + +A token bucket implementation for rate limiting in Nim. + +## Overview + +This library provides an enhanced token bucket implementation that addresses issues found in the original `chronos/rate_limit.nim` and adds advanced features for better rate limiting control. The token bucket algorithm is commonly used for rate limiting in network applications, APIs, and distributed systems. + +## Background + +This is an extract from `chronos/rate_limit.nim` due to a bug in the original implementation that cannot be solved without harming the original features of the TokenBucket class. This implementation serves as a shortcut to enable progress with the nwaku rate limiter implementation. + +**Reference**: [nim-chronos issue #500](https://github.com/status-im/nim-chronos/issues/500) + +## Key Features + +This version differs from the original `chronos/rate_limit.nim` in several ways: + +- **Compensating Mode**: Calculates compensation as unused bucket capacity from previous periods (up to 25% threshold), with proper time period calculation to avoid overcompensation during non-usage periods +- **Strict Mode**: Replenishes only when the time period is over, filling the bucket to maximum capacity + +## Installation + +Add this to your `.nimble` file: + +```nim +requires "https://github.com/waku-org/token_bucket" +``` + +Or install directly: + +```bash +nimble install https://github.com/waku-org/token_bucket +``` diff --git a/tests/test_token_bucket.nim b/tests/test_token_bucket.nim new file mode 100644 index 0000000..38f79ba --- /dev/null +++ b/tests/test_token_bucket.nim @@ -0,0 +1,251 @@ +# Chronos Test Suite +# (c) Copyright 2022-Present +# Status Research & Development GmbH +# +# Licensed under either of +# Apache License, version 2.0, (LICENSE-APACHEv2) +# MIT license (LICENSE-MIT) + +{.used.} + +import testutils/unittests +import chronos +import ../token_bucket/token_bucket + +suite "Token Bucket": + test "TokenBucket Sync test - strict": + var bucket = TokenBucket.newStrict(1000, 1.milliseconds) + let + start = Moment.now() + fullTime = start + 1.milliseconds + check: + bucket.tryConsume(800, start) == true + bucket.tryConsume(200, start) == true + # Out of budget + bucket.tryConsume(100, start) == false + bucket.tryConsume(800, fullTime) == true + bucket.tryConsume(200, fullTime) == true + # Out of budget + bucket.tryConsume(100, fullTime) == false + + test "TokenBucket Sync test - compensating": + var bucket = TokenBucket.new(1000, 1.milliseconds) + let + start = Moment.now() + fullTime = start + 1.milliseconds + check: + bucket.tryConsume(800, start) == true + bucket.tryConsume(200, start) == true + # Out of budget + bucket.tryConsume(100, start) == false + bucket.tryConsume(800, fullTime) == true + bucket.tryConsume(200, fullTime) == true + # Due not using the bucket for a full period the compensation will satisfy this request + bucket.tryConsume(100, fullTime) == true + + test "TokenBucket Max compensation": + var bucket = TokenBucket.new(1000, 1.minutes) + var reqTime = Moment.now() + + check bucket.tryConsume(1000, reqTime) + check bucket.tryConsume(1, reqTime) == false + reqTime += 1.minutes + check bucket.tryConsume(500, reqTime) == true + reqTime += 1.minutes + check bucket.tryConsume(1000, reqTime) == true + reqTime += 10.seconds + # max compensation is 25% so try to consume 250 more + check bucket.tryConsume(250, reqTime) == true + reqTime += 49.seconds + # out of budget within the same period + check bucket.tryConsume(1, reqTime) == false + + test "TokenBucket Short replenish": + 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 + + test "TokenBucket getAvailableCapacity strict": + let now = Moment.now() + var reqTime = now + var bucket = + TokenBucket.new(1000, 1.minutes, ReplenishMode.Strict, lastTimeFull = now) + + # Test full bucket capacity ratio + let lastTimeFull_1 = reqTime + check bucket.getAvailableCapacity(reqTime) == (1000, 1000, lastTimeFull_1) + + # # 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, lastTimeFull_2) + + # Consume more tokens + reqTime += 1.seconds + check bucket.tryConsume(300, reqTime) == true + check bucket.getAvailableCapacity(reqTime) == (300, 1000, lastTimeFull_2) + + # Test when period has elapsed (should return 1.0) + reqTime += 1.minutes + 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, lastTimeFull_3) + + test "TokenBucket getAvailableCapacity compensating": + 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, 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, lastTimeFull_2) + + # Consume more tokens + reqTime += 1.seconds + check bucket.tryConsume(300, reqTime) == true + 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) + 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 + + 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 diff --git a/token_bucket.nimble b/token_bucket.nimble new file mode 100644 index 0000000..9d248a7 --- /dev/null +++ b/token_bucket.nimble @@ -0,0 +1,9 @@ +# Package + +version = "0.0.1" +author = "Nwaku Team" +description = "Token bucket for rate limiting" +license = "MIT" + +### Dependencies +requires "nim >= 2.2.4", "chronos" diff --git a/token_bucket/token_bucket.nim b/token_bucket/token_bucket.nim new file mode 100644 index 0000000..f52cc8a --- /dev/null +++ b/token_bucket/token_bucket.nim @@ -0,0 +1,204 @@ +{.push raises: [].} + +import chronos, std/math, std/options + +const BUDGET_COMPENSATION_LIMIT_PERCENT = 0.25 + +## This is an extract from chronos/rate_limit.nim due to the found bug in the original implementation. +## Unfortunately that bug cannot be solved without harm the original features of TokenBucket class. +## So, this current shortcut is used to enable move ahead with nwaku rate limiter implementation. +## ref: https://github.com/status-im/nim-chronos/issues/500 +## +## This version of TokenBucket is different from the original one in chronos/rate_limit.nim in many ways: +## - It has a new mode called `Compensating` which is the default mode. +## Compensation is calculated as the not used bucket capacity in the last measured period(s) in average. +## or up until maximum the allowed compansation treshold (Currently it is const 25%). +## Also compensation takes care of the proper time period calculation to avoid non-usage periods that can lead to +## overcompensation. +## - Strict mode is also available which will only replenish when time period is over but also will fill +## the bucket to the max capacity. + +type + ReplenishMode* = enum + Strict + Compensating + + TokenBucket* = ref object + budget: int ## Current number of tokens in the bucket + budgetCap: int ## Bucket capacity + lastTimeFull: Moment + ## This timer measures the proper periodizaiton of the bucket refilling + fillDuration: Duration ## Refill period + case replenishMode*: ReplenishMode + of Strict: + ## In strict mode, the bucket is refilled only till the budgetCap + discard + of Compensating: + ## This is the default mode. + maxCompensation: float + +func periodDistance(bucket: TokenBucket, currentTime: Moment): float = + ## notice fillDuration cannot be zero by design + ## period distance is a float number representing the calculated period time + ## since the last time bucket was refilled. + return + nanoseconds(currentTime - bucket.lastTimeFull).float / + nanoseconds(bucket.fillDuration).float + +func getUsageAverageSince(bucket: TokenBucket, distance: float): float = + if distance == 0.float: + ## in case there is zero time difference than the usage percentage is 100% + return 1.0 + + ## budgetCap can never be zero + ## usage average is calculated as a percentage of total capacity available over + ## the measured period + return bucket.budget.float / bucket.budgetCap.float / distance + +proc calcCompensation(bucket: TokenBucket, averageUsage: float): int = + # if we already fully used or even overused the tokens, there is no place for compensation + if averageUsage >= 1.0: + return 0 + + ## compensation is the not used bucket capacity in the last measured period(s) in average. + ## or maximum the allowed compansation treshold + let compensationPercent = + min((1.0 - averageUsage) * bucket.budgetCap.float, bucket.maxCompensation) + return trunc(compensationPercent).int + +func periodElapsed(bucket: TokenBucket, currentTime: Moment): bool = + return currentTime - bucket.lastTimeFull >= bucket.fillDuration + +## Update will take place if bucket is empty and trying to consume tokens. +## It checks if the bucket can be replenished as refill duration is passed or not. +## - strict mode: +proc updateStrict(bucket: TokenBucket, currentTime: Moment) = + if bucket.fillDuration == default(Duration): + bucket.budget = min(bucket.budgetCap, bucket.budget) + return + + if not periodElapsed(bucket, currentTime): + return + + bucket.budget = bucket.budgetCap + bucket.lastTimeFull = currentTime + +## - compensating - ballancing load: +## - between updates we calculate average load (current bucket capacity / number of periods till last update) +## - gives the percentage load used recently +## - with this we can replenish bucket up to 100% + calculated leftover from previous period (caped with max treshold) +proc updateWithCompensation(bucket: TokenBucket, currentTime: Moment) = + if bucket.fillDuration == default(Duration): + bucket.budget = min(bucket.budgetCap, bucket.budget) + return + + # do not replenish within the same period + if not periodElapsed(bucket, currentTime): + return + + let distance = bucket.periodDistance(currentTime) + let recentAvgUsage = bucket.getUsageAverageSince(distance) + let compensation = bucket.calcCompensation(recentAvgUsage) + + bucket.budget = bucket.budgetCap + compensation + bucket.lastTimeFull = currentTime + +proc update(bucket: TokenBucket, currentTime: Moment) = + if bucket.replenishMode == ReplenishMode.Compensating: + updateWithCompensation(bucket, currentTime) + else: + updateStrict(bucket, currentTime) + +proc getAvailableCapacity*( + bucket: TokenBucket, currentTime: Moment = Moment.now() +): 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, 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, 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, + ## Otherwhise, return false. + + if bucket.budget >= bucket.budgetCap: + bucket.lastTimeFull = now + + if bucket.budget >= tokens: + bucket.budget -= tokens + return true + + bucket.update(now) + + if bucket.budget >= tokens: + bucket.budget -= tokens + return true + else: + return false + +proc replenish*(bucket: TokenBucket, tokens: int, now = Moment.now()) = + ## Add `tokens` to the budget (capped to the bucket capacity) + bucket.budget += tokens + bucket.update(now) + +proc new*( + T: type[TokenBucket], + 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: actualBudget, + budgetCap: budgetCap, + fillDuration: fillDuration, + lastTimeFull: lastTimeFull, + replenishMode: mode, + ) + of ReplenishMode.Compensating: + T( + budget: actualBudget, + budgetCap: budgetCap, + fillDuration: fillDuration, + lastTimeFull: lastTimeFull, + replenishMode: mode, + maxCompensation: budgetCap.float * BUDGET_COMPENSATION_LIMIT_PERCENT, + ) + +proc newStrict*(T: type[TokenBucket], capacity: int, period: Duration): TokenBucket = + T.new(capacity, period, ReplenishMode.Strict) + +proc newCompensating*( + T: type[TokenBucket], capacity: int, period: Duration +): TokenBucket = + T.new(capacity, period, ReplenishMode.Compensating) + +func `$`*(b: TokenBucket): string {.inline.} = + if isNil(b): + return "nil" + return $b.budgetCap & "/" & $b.fillDuration + +func `$`*(ob: Option[TokenBucket]): string {.inline.} = + if ob.isNone(): + return "no-limit" + + return $ob.get()