Merge 4e49dc26d2f64cb16d617fda4e29b8393344f2e5 into dafdee9f5ffc0460f45307c61fbd8e9832fc3ecd

tests/common/test_tokenbucket.nim

@@ -67,3 +67,185 @@ suite "Token Bucket":
     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

waku/common/rate_limit/token_bucket.nim

@@ -109,6 +109,23 @@ proc update(bucket: TokenBucket, currentTime: Moment) =
   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.
@@ -138,26 +155,31 @@ proc new*(
     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: budgetCap,
+      budget: actualBudget,
       budgetCap: budgetCap,
       fillDuration: fillDuration,
-      lastTimeFull: Moment.now(),
+      lastTimeFull: lastTimeFull,
       replenishMode: mode,
     )
   of ReplenishMode.Compensating:
     T(
-      budget: budgetCap,
+      budget: actualBudget,
       budgetCap: budgetCap,
       fillDuration: fillDuration,
-      lastTimeFull: Moment.now(),
+      lastTimeFull: lastTimeFull,
       replenishMode: mode,
       maxCompensation: budgetCap.float * BUDGET_COMPENSATION_LIMIT_PERCENT,
     )