chore: token bucket unification

tests/common/test_all.nim

@@ -8,7 +8,6 @@ import
   ./test_protobuf_validation,
   ./test_sqlite_migrations,
   ./test_parse_size,
-  ./test_tokenbucket,
   ./test_requestratelimiter,
   ./test_ratelimit_setting,
   ./test_timed_map

tests/common/test_requestratelimiter.nim

@@ -44,8 +44,9 @@ suite "RequestRateLimiter":
     # conn1 reached the 75% of the main bucket over 2 periods of time
     check limiter.checkUsage(proto, conn1, now + 3.minutes) == false
 
-    # conn2 has not used its tokens while we have 1 more tokens left in the main bucket
-    check limiter.checkUsage(proto, conn2, now + 3.minutes) == true
+    # conn2 has not used its tokens but while we have 0 more tokens left in the main bucket
+    # it will tolerate some burst within period.
+    check limiter.checkUsage(proto, conn2, now + 3.minutes + 50.seconds) == true
 
   test "RequestRateLimiter Restrict overusing peer":
     # keep limits low for easier calculation of ratios
@@ -55,6 +56,8 @@ suite "RequestRateLimiter":
     # as ratio is 2 in this case but max tokens are main tokens*ratio . 0.75
     # notice meanwhile we have 20 tokens over 2 period (4 mins) in sum
     # See: waku/common/rate_limit/request_limiter.nim #func calcPeriodRatio
+    # This tests shows balanced token utilization as being replenished in between periods but never overusing
+    # intended rates.
 
     let now = Moment.now()
     # with first use we register the peer also and start its timer
@@ -64,23 +67,22 @@ suite "RequestRateLimiter":
     # run out of main tokens but still used one more token from the peer's bucket
     check limiter.checkUsage(proto, conn1, now) == false
 
-    for i in 0 ..< 4:
+    for i in 0 ..< 10:
       check limiter.checkUsage(proto, conn1, now + 3.minutes) == true
 
     # conn1 reached the 75% of the main bucket over 2 periods of time
     check limiter.checkUsage(proto, conn1, now + 3.minutes) == false
 
-    check limiter.checkUsage(proto, conn2, now + 3.minutes) == true
-    check limiter.checkUsage(proto, conn2, now + 3.minutes) == true
-    check limiter.checkUsage(proto, conn3, now + 3.minutes) == true
-    check limiter.checkUsage(proto, conn2, now + 3.minutes) == true
-    check limiter.checkUsage(proto, conn3, now + 3.minutes) == true
+    check limiter.checkUsage(proto, conn2, now + 3.minutes + 50.seconds) == true
+    check limiter.checkUsage(proto, conn2, now + 3.minutes + 50.seconds) == true
+    check limiter.checkUsage(proto, conn3, now + 3.minutes + 50.seconds) == true
+    check limiter.checkUsage(proto, conn2, now + 3.minutes + 50.seconds) == true
 
     # conn1 gets replenished as the ratio was 2 giving twice as long replenish period than the main bucket
     # see waku/common/rate_limit/request_limiter.nim #func calcPeriodRatio and calcPeerTokenSetting
     check limiter.checkUsage(proto, conn1, now + 4.minutes) == true
     # requests of other peers can also go
-    check limiter.checkUsage(proto, conn2, now + 4100.milliseconds) == true
+    check limiter.checkUsage(proto, conn2, now + 5.minutes) == true
     check limiter.checkUsage(proto, conn3, now + 5.minutes) == true
 
   test "RequestRateLimiter lowest possible volume":

waku/common/rate_limit/per_peer_limiter.nim

@@ -6,11 +6,10 @@
 
 {.push raises: [].}
 
-import std/[options, tables], libp2p/stream/connection
-
+import std/[options, tables], libp2p/stream/connection, chronos/ratelimit
 import ./[single_token_limiter, service_metrics], ../../utils/tableutils
 
-export token_bucket, setting, service_metrics
+export setting, service_metrics
 
 type PerPeerRateLimiter* = ref object of RootObj
   setting*: Option[RateLimitSetting]
@@ -20,7 +19,7 @@ proc mgetOrPut(
     perPeerRateLimiter: var PerPeerRateLimiter, peerId: PeerId
 ): var Option[TokenBucket] =
   return perPeerRateLimiter.peerBucket.mgetOrPut(
-    peerId, newTokenBucket(perPeerRateLimiter.setting, ReplenishMode.Compensating)
+    peerId, newTokenBucket(perPeerRateLimiter.setting, ReplenishMode.Balanced)
   )
 
 template checkUsageLimit*(

waku/common/rate_limit/request_limiter.nim

@@ -19,14 +19,15 @@ import
   std/[options, math],
   chronicles,
   chronos/timer,
+  chronos/ratelimit,
   libp2p/stream/connection,
   libp2p/utility
 
 import std/times except TimeInterval, Duration, seconds, minutes
 
 import ./[single_token_limiter, service_metrics, timed_map]
 
-export token_bucket, setting, service_metrics
+export setting, service_metrics
 
 logScope:
   topics = "waku ratelimit"
@@ -59,6 +60,9 @@ proc checkUsage*(
     return true
 
   let peerBucket = t.mgetOrPut(conn.peerId)
+  let avail = peerBucket.getAvailableCapacity(now)
+  let globAvail = t.tokenBucket.get().getAvailableCapacity(now)
+
   ## check requesting peer's usage is not over the calculated ratio and let that peer go which not requested much/or this time...
   if not peerBucket.tryConsume(1, now):
     trace "peer usage limit reached", peer = conn.peerId

waku/common/rate_limit/single_token_limiter.nim

@@ -2,24 +2,24 @@
 
 {.push raises: [].}
 
-import std/[options], chronos/timer, libp2p/stream/connection, libp2p/utility
+import std/[options], chronos/timer, chronos/ratelimit, libp2p/stream/connection, libp2p/utility
 
 import std/times except TimeInterval, Duration
 
-import ./[token_bucket, setting, service_metrics]
-export token_bucket, setting, service_metrics
+import ./[setting, service_metrics]
+export setting, service_metrics
 
 proc newTokenBucket*(
     setting: Option[RateLimitSetting],
-    replenishMode: ReplenishMode = ReplenishMode.Compensating,
+    replenishMode: ReplenishMode = ReplenishMode.Balanced,
 ): Option[TokenBucket] =
   if setting.isNone():
     return none[TokenBucket]()
 
   if setting.get().isUnlimited():
     return none[TokenBucket]()
 
-  return some(TokenBucket.new(setting.get().volume, setting.get().period))
+  return some(TokenBucket.new(setting.get().volume, setting.get().period, replenishMode))
 
 proc checkUsage(
     t: var TokenBucket, proto: string, now = Moment.now()