Less flaky rate limit tests (#408)

chronos/ratelimit.nim

@@ -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*(

tests/testratelimit.nim

@@ -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(800, start) == true
-      bucket.tryConsume(200) == true
+      bucket.tryConsume(200, start) == true
 
       # Out of budget
-      bucket.tryConsume(100) == false
+      bucket.tryConsume(100, start) == false
-    waitFor(sleepAsync(10.milliseconds))
+      bucket.tryConsume(800, fullTime) == true
-    check:
+      bucket.tryConsume(200, fullTime) == true
-      bucket.tryConsume(800) == true
-      bucket.tryConsume(200) == 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))
+    waitFor(bucket.consume(1000).wait(5.seconds))
-    when not defined(macosx):
+    check Moment.now() - beforeStart in 900.milliseconds .. 1500.milliseconds
-      # 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
 
   test "Sync manual replenish":
     var bucket = TokenBucket.new(1000, 0.seconds)
+    let start = Moment.now()
     check:
-      bucket.tryConsume(1000) == true
+      bucket.tryConsume(1000, start) == true
-      bucket.tryConsume(1000) == false
+      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)
+    let start = Moment.now()
-    check bucket.tryConsume(1) == false
+    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()
+    var fakeNow = start
-    while Moment.now() - start >= 514.milliseconds:
+    while fakeNow - start < 514.milliseconds:
-      check bucket.tryConsume(1) == false
+      check bucket.tryConsume(1, fakeNow) == false
-      waitFor(sleepAsync(10.milliseconds))
+      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)
+    let start = Moment.now()
-    check bucket.tryConsume(1) == false
+    check bucket.tryConsume(15000, start)
+    check bucket.tryConsume(1, start) == false
 
-    waitFor(sleepAsync(1.milliseconds))
+    check bucket.tryConsume(15000, start + 1.milliseconds) == true
-    check bucket.tryConsume(15000) == true