mirror of
https://github.com/logos-messaging/token_bucket.git
synced 2026-01-02 14:13:08 +00:00
initial
This commit is contained in:
pablo
commit
fe54bb26d3
24
.gitignore
vendored
Normal file
24
.gitignore
vendored
Normal file
@ -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
|
||||
56
Makefile
Normal file
56
Makefile
Normal file
@ -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
|
||||
34
README.md
Normal file
34
README.md
Normal file
@ -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
|
||||
```
|
||||
251
tests/test_token_bucket.nim
Normal file
251
tests/test_token_bucket.nim
Normal file
@ -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
|
||||
9
token_bucket.nimble
Normal file
9
token_bucket.nimble
Normal file
@ -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"
|
||||
204
token_bucket/token_bucket.nim
Normal file
204
token_bucket/token_bucket.nim
Normal file
@ -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()
|
||||
Loading…
x
Reference in New Issue
Block a user