codex-storage
/
nim-chronos
mirror of https://github.com/codex-storage/nim-chronos.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Less flaky rate limit tests (#408)

This commit is contained in:
Tanguy 2023-06-23 10:11:14 +02:00 committed by GitHub
parent 0035f4fa66
commit 47016961f5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 41 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
chronos/ratelimit.nim
View File

@ -28,13 +28,15 @@ type
pendingRequests: seq[BucketWaiter]
manuallyReplenished: AsyncEvent
proc update(bucket: TokenBucket) =
proc update(bucket: TokenBucket, currentTime: Moment) =
if bucket.fillDuration == default(Duration):
bucket.budget = min(bucket.budgetCap, bucket.budget)
return
if currentTime < bucket.lastUpdate:
return
let
currentTime = Moment.now()
timeDelta = currentTime - bucket.lastUpdate
fillPercent = timeDelta.milliseconds.float / bucket.fillDuration.milliseconds.float
replenished =
@ -46,7 +48,7 @@ proc update(bucket: TokenBucket) =
bucket.lastUpdate += milliseconds(deltaFromReplenished)
bucket.budget = min(bucket.budgetCap, bucket.budget + replenished)
proc tryConsume*(bucket: TokenBucket, tokens: int): bool =
proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool =
## If `tokens` are available, consume them,
## Otherwhise, return false.
@ -54,7 +56,7 @@ proc tryConsume*(bucket: TokenBucket, tokens: int): bool =
bucket.budget -= tokens
return true
bucket.update()
bucket.update(now)
if bucket.budget >= tokens:
bucket.budget -= tokens
@ -93,12 +95,12 @@ proc worker(bucket: TokenBucket) {.async.} =
bucket.workFuture = nil
proc consume*(bucket: TokenBucket, tokens: int): Future[void] =
proc consume*(bucket: TokenBucket, tokens: int, now = Moment.now()): Future[void] =
## Wait for `tokens` to be available, and consume them.
let retFuture = newFuture[void]("TokenBucket.consume")
if isNil(bucket.workFuture) or bucket.workFuture.finished():
if bucket.tryConsume(tokens):
if bucket.tryConsume(tokens, now):
retFuture.complete()
return retFuture
@ -119,10 +121,10 @@ proc consume*(bucket: TokenBucket, tokens: int): Future[void] =
return retFuture
proc replenish*(bucket: TokenBucket, tokens: int) =
proc replenish*(bucket: TokenBucket, tokens: int, now = Moment.now()) =
## Add `tokens` to the budget (capped to the bucket capacity)
bucket.budget += tokens
bucket.update()
bucket.update(now)
bucket.manuallyReplenished.fire()
proc new*(

62
tests/testratelimit.nim
View File

@ -15,22 +15,23 @@ import ../chronos/ratelimit
suite "Token Bucket":
test "Sync test":
var bucket = TokenBucket.new(1000, 1.milliseconds)
let
start = Moment.now()
fullTime = start + 1.milliseconds
check:
bucket.tryConsume(800) == true
bucket.tryConsume(200) == true
bucket.tryConsume(800, start) == true
bucket.tryConsume(200, start) == true
# Out of budget
bucket.tryConsume(100) == false
waitFor(sleepAsync(10.milliseconds))
check:
bucket.tryConsume(800) == true
bucket.tryConsume(200) == true
bucket.tryConsume(100, start) == false
bucket.tryConsume(800, fullTime) == true
bucket.tryConsume(200, fullTime) == true
# Out of budget
bucket.tryConsume(100) == false
bucket.tryConsume(100, fullTime) == false
test "Async test":
var bucket = TokenBucket.new(1000, 500.milliseconds)
var bucket = TokenBucket.new(1000, 1000.milliseconds)
check: bucket.tryConsume(1000) == true
var toWait = newSeq[Future[void]]()
@ -41,28 +42,26 @@ suite "Token Bucket":
waitFor(allFutures(toWait))
let duration = Moment.now() - start
check: duration in 700.milliseconds .. 1100.milliseconds
check: duration in 1400.milliseconds .. 2200.milliseconds
test "Over budget async":
var bucket = TokenBucket.new(100, 10.milliseconds)
var bucket = TokenBucket.new(100, 100.milliseconds)
# Consume 10* the budget cap
let beforeStart = Moment.now()
waitFor(bucket.consume(1000).wait(1.seconds))
when not defined(macosx):
# CI's macos scheduler is so jittery that this tests sometimes takes >500ms
# the test will still fail if it's >1 seconds
check Moment.now() - beforeStart in 90.milliseconds .. 150.milliseconds
waitFor(bucket.consume(1000).wait(5.seconds))
check Moment.now() - beforeStart in 900.milliseconds .. 1500.milliseconds
test "Sync manual replenish":
var bucket = TokenBucket.new(1000, 0.seconds)
let start = Moment.now()
check:
bucket.tryConsume(1000) == true
bucket.tryConsume(1000) == false
bucket.tryConsume(1000, start) == true
bucket.tryConsume(1000, start) == false
bucket.replenish(2000)
check:
bucket.tryConsume(1000) == true
bucket.tryConsume(1000, start) == true
# replenish is capped to the bucket max
bucket.tryConsume(1000) == false
bucket.tryConsume(1000, start) == false
test "Async manual replenish":
var bucket = TokenBucket.new(10 * 150, 0.seconds)
@ -102,24 +101,25 @@ suite "Token Bucket":
test "Very long replenish":
var bucket = TokenBucket.new(7000, 1.hours)
check bucket.tryConsume(7000)
check bucket.tryConsume(1) == false
let start = Moment.now()
check bucket.tryConsume(7000, start)
check bucket.tryConsume(1, start) == false
# With this setting, it takes 514 milliseconds
# to tick one. Check that we can eventually
# consume, even if we update multiple time
# before that
let start = Moment.now()
while Moment.now() - start >= 514.milliseconds:
check bucket.tryConsume(1) == false
waitFor(sleepAsync(10.milliseconds))
var fakeNow = start
while fakeNow - start < 514.milliseconds:
check bucket.tryConsume(1, fakeNow) == false
fakeNow += 30.milliseconds
check bucket.tryConsume(1) == false
check bucket.tryConsume(1, fakeNow) == true
test "Short replenish":
var bucket = TokenBucket.new(15000, 1.milliseconds)
check bucket.tryConsume(15000)
check bucket.tryConsume(1) == false
let start = Moment.now()
check bucket.tryConsume(15000, start)
check bucket.tryConsume(1, start) == false
waitFor(sleepAsync(1.milliseconds))
check bucket.tryConsume(15000) == true
check bucket.tryConsume(15000, start + 1.milliseconds) == true