From 7c8b6ed5de8fd672f9e00820ba956306f73c9d7e Mon Sep 17 00:00:00 2001
From: Mykola Zyk <mykola.zyk@dinecat.com>
Date: Wed, 22 Dec 2021 20:01:14 +0200
Subject: [PATCH] [RateLimiter] TokenBucket policy fix for adding tokens with a
 predefined frequency

---
 .../Component/RateLimiter/Policy/Rate.php     | 12 +++++++++
 .../RateLimiter/Policy/TokenBucket.php        |  7 ++++-
 .../RateLimiter/Policy/TokenBucketLimiter.php |  2 --
 .../Tests/Policy/TokenBucketLimiterTest.php   | 26 +++++++++++++++++++
 4 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/src/Symfony/Component/RateLimiter/Policy/Rate.php b/src/Symfony/Component/RateLimiter/Policy/Rate.php
index 0c91ef78e76c2..3775f53ae8b3c 100644
--- 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
index e4eb32a744a71..2d43286e15e23 100644
--- 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
index 09c4e36cdf861..5724eba2b2abb 100644
--- 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
index 84136ed7f5d7d..3b7b579c0cf77 100644
--- 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);