fix: changed ratio to actual capacity and PR comments

tests/common/test_tokenbucket.nim

@@ -10,13 +10,8 @@
 
 import testutils/unittests
 import chronos
-import std/math
 import ../../waku/common/rate_limit/token_bucket
 
-# Helper function for approximate float equality
-proc equals(a, b: float, tolerance: float = 1e-6): bool =
-  abs(a - b) <= tolerance
-
 suite "Token Bucket":
   test "TokenBucket Sync test - strict":
     var bucket = TokenBucket.newStrict(1000, 1.milliseconds)
@@ -73,27 +68,76 @@ suite "Token Bucket":
 
     check bucket.tryConsume(15000, start + 1.milliseconds) == true
 
-  test "TokenBucket getAvailableCapacityRatio":
+  test "TokenBucket getAvailableCapacity strict":
     var bucket = TokenBucket.new(1000, 1.minutes, ReplenishMode.Strict)
     var reqTime = Moment.now()
 
     # Test full bucket capacity ratio
-    check equals(bucket.getAvailableCapacityRatio(reqTime), 1.0) # 1000/1000 = 1.0
+    check bucket.getAvailableCapacity(reqTime) == (1000, 1000)
 
     # Consume some tokens and check ratio
     reqTime += 1.seconds
     check bucket.tryConsume(400, reqTime) == true
-    check equals(bucket.getAvailableCapacityRatio(reqTime), 0.6) # 600/1000 = 0.6
+    check bucket.getAvailableCapacity(reqTime) == (600, 1000)
 
     # Consume more tokens
     reqTime += 1.seconds
     check bucket.tryConsume(300, reqTime) == true
-    check equals(bucket.getAvailableCapacityRatio(reqTime), 0.3) # 300/1000 = 0.3
+    check bucket.getAvailableCapacity(reqTime) == (300, 1000)
 
     # Test when period has elapsed (should return 1.0)
     reqTime += 1.minutes
-    check equals(bucket.getAvailableCapacityRatio(reqTime), 1.0) # 1000/1000 = 1.0
+    check bucket.getAvailableCapacity(reqTime) == (1000, 1000)
 
     # Test with empty bucket
     check bucket.tryConsume(1000, reqTime) == true
-    check equals(bucket.getAvailableCapacityRatio(reqTime), 0.0) # 0/1000 = 0.0
+    check bucket.getAvailableCapacity(reqTime) == (0, 1000)
+
+  test "TokenBucket getAvailableCapacity compensating":
+    var bucket = TokenBucket.new(1000, 1.minutes, ReplenishMode.Compensating)
+    var reqTime = Moment.now()
+
+    # Test full bucket capacity
+    check bucket.getAvailableCapacity(reqTime) == (1000, 1000)
+
+    # Consume some tokens and check available capacity
+    reqTime += 1.seconds
+    check bucket.tryConsume(400, reqTime) == true
+    check bucket.getAvailableCapacity(reqTime) == (600, 1000)
+
+    # Consume more tokens
+    reqTime += 1.seconds
+    check bucket.tryConsume(300, reqTime) == true
+    check bucket.getAvailableCapacity(reqTime) == (300, 1000)
+
+    # 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
+

waku/common/rate_limit/token_bucket.nim

@@ -109,11 +109,19 @@ proc update(bucket: TokenBucket, currentTime: Moment) =
   else:
     updateStrict(bucket, currentTime)
 
-## Returns the available capacity ratio of the bucket: 0.0 (empty); 1.0 (full)
+## Returns the available capacity of the bucket: (budget, budgetCap)
-proc getAvailableCapacityRatio*(bucket: TokenBucket, currentTime: Moment): float =
+proc getAvailableCapacity*(bucket: TokenBucket, currentTime: Moment = Moment.now()): tuple[budget: int, budgetCap: int] =
   if periodElapsed(bucket, currentTime):
-    return 1.0
+    case bucket.replenishMode
-  return bucket.budget.float / bucket.budgetCap.float
+    of ReplenishMode.Strict:
+      return (bucket.budgetCap, bucket.budgetCap)
+    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)
+  return (bucket.budget, bucket.budgetCap)
 
 proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool =
   ## If `tokens` are available, consume them,