[RateLimiter] TokenBucket policy fix for adding tokens with a predefined frequency

relo-san · fabpot · commit 7c8b6ed5de8f · 2023-10-10T07:34:14.000+02:00
diff --git a/src/Symfony/Component/RateLimiter/Policy/Rate.php b/src/Symfony/Component/RateLimiter/Policy/Rate.php
@@ -95,6 +95,18 @@ public function calculateNewTokensDuringInterval(float $duration): int
         return $cycles * $this->refillAmount;
     }
 
+    /**
+     * Calculates total amount in seconds of refill intervals during $duration (for maintain strict refill frequency).
+     *
+     * @param float $duration interval in seconds
+     */
+    public function calculateRefillInterval(float $duration): int
+    {
+        $cycleTime = TimeUtil::dateIntervalToSeconds($this->refillTime);
+
+        return floor($duration / $cycleTime) * $cycleTime;
+    }
+
     public function __toString(): string
     {
         return $this->refillTime->format('P%dDT%HH%iM%sS').'-'.$this->refillAmount;
diff --git a/src/Symfony/Component/RateLimiter/Policy/TokenBucket.php b/src/Symfony/Component/RateLimiter/Policy/TokenBucket.php
@@ -83,8 +83,13 @@ public function setTokens(int $tokens): void
     public function getAvailableTokens(float $now): int
     {
         $elapsed = max(0, $now - $this->timer);
+        $newTokens = $this->rate->calculateNewTokensDuringInterval($elapsed);
 
-        return min($this->burstSize, $this->tokens + $this->rate->calculateNewTokensDuringInterval($elapsed));
+        if ($newTokens > 0) {
+            $this->timer += $this->rate->calculateRefillInterval($elapsed);
+        }
+
+        return min($this->burstSize, $this->tokens + $newTokens);
     }
 
     public function getExpirationTime(): int
diff --git a/src/Symfony/Component/RateLimiter/Policy/TokenBucketLimiter.php b/src/Symfony/Component/RateLimiter/Policy/TokenBucketLimiter.php
@@ -72,7 +72,6 @@ public function reserve(int $tokens = 1, float $maxTime = null): Reservation
             if ($availableTokens >= $tokens) {
                 // tokens are now available, update bucket
                 $bucket->setTokens($availableTokens - $tokens);
-                $bucket->setTimer($now);
 
                 $reservation = new Reservation($now, new RateLimit($bucket->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now)), true, $this->maxBurst));
             } else {
@@ -89,7 +88,6 @@ public function reserve(int $tokens = 1, float $maxTime = null): Reservation
                 // at $now + $waitDuration all tokens will be reserved for this process,
                 // so no tokens are left for other processes.
                 $bucket->setTokens($availableTokens - $tokens);
-                $bucket->setTimer($now);
 
                 $reservation = new Reservation($now + $waitDuration, new RateLimit(0, \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false, $this->maxBurst));
             }
diff --git a/src/Symfony/Component/RateLimiter/Tests/Policy/TokenBucketLimiterTest.php b/src/Symfony/Component/RateLimiter/Tests/Policy/TokenBucketLimiterTest.php
@@ -128,6 +128,32 @@ public function testBucketResilientToTimeShifting()
         $this->assertSame(100, $bucket->getAvailableTokens($serverOneClock));
     }
 
+    public function testBucketRefilledWithStrictFrequency()
+    {
+        $limiter = $this->createLimiter(1000, new Rate(\DateInterval::createFromDateString('15 seconds'), 100));
+        $rateLimit = $limiter->consume(300);
+
+        $this->assertTrue($rateLimit->isAccepted());
+        $this->assertEquals(700, $rateLimit->getRemainingTokens());
+
+        $expected = 699;
+
+        for ($i = 1; $i <= 20; ++$i) {
+            $rateLimit = $limiter->consume();
+            $this->assertTrue($rateLimit->isAccepted());
+            $this->assertEquals($expected, $rateLimit->getRemainingTokens());
+
+            sleep(4);
+            --$expected;
+
+            if (\in_array($i, [4, 8, 12], true)) {
+                $expected += 100;
+            } elseif (\in_array($i, [15, 19], true)) {
+                $expected = 999;
+            }
+        }
+    }
+
     private function createLimiter($initialTokens = 10, Rate $rate = null)
     {
         return new TokenBucketLimiter('test', $initialTokens, $rate ?? Rate::perSecond(10), $this->storage);

Original file line number	Diff line number	Diff line change
`@@ -83,8 +83,13 @@ public function setTokens(int $tokens): void`
`83`	`83`	`public function getAvailableTokens(float $now): int`
`84`	`84`	`{`
`85`	`85`	`$elapsed = max(0, $now - $this->timer);`
	`86`	`+ $newTokens = $this->rate->calculateNewTokensDuringInterval($elapsed);`
`86`	`87`
`87`		`- return min($this->burstSize, $this->tokens + $this->rate->calculateNewTokensDuringInterval($elapsed));`
	`88`	`+ if ($newTokens > 0) {`
	`89`	`+ $this->timer += $this->rate->calculateRefillInterval($elapsed);`
	`90`	`+ }`
	`91`	`+`
	`92`	`+ return min($this->burstSize, $this->tokens + $newTokens);`
`88`	`93`	`}`
`89`	`94`
`90`	`95`	`public function getExpirationTime(): int`