mirror of
https://github.com/logos-messaging/token_bucket.git
synced 2026-01-02 14:13:08 +00:00
252 lines
9.8 KiB
Nim
252 lines
9.8 KiB
Nim
# 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 ../src/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
|