Make RateLimiter resilient to timeShifting

symfony · Dec 2, 2021 · 7eec0d6 · 7eec0d6
1 parent 2183551
commit 7eec0d6
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 11 deletions.
diff --git a/Policy/FixedWindowLimiter.php b/Policy/FixedWindowLimiter.php
@@ -66,7 +66,7 @@ public function reserve(int $tokens = 1, float $maxTime = null): Reservation
             $now = microtime(true);
             $availableTokens = $window->getAvailableTokens($now);
             if ($availableTokens >= $tokens) {
-                $window->add($tokens);
+                $window->add($tokens, $now);
 
                 $reservation = new Reservation($now, new RateLimit($window->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now)), true, $this->limit));
             } else {
@@ -77,7 +77,7 @@ public function reserve(int $tokens = 1, float $maxTime = null): Reservation
                     throw new MaxWaitDurationExceededException(sprintf('The rate limiter wait time ("%d" seconds) is longer than the provided maximum time ("%d" seconds).', $waitDuration, $maxTime), new RateLimit($window->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false, $this->limit));
                 }
 
-                $window->add($tokens);
+                $window->add($tokens, $now);
 
                 $reservation = new Reservation($now + $waitDuration, new RateLimit($window->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false, $this->limit));
             }

diff --git a/Policy/TokenBucket.php b/Policy/TokenBucket.php
@@ -82,7 +82,7 @@ public function setTokens(int $tokens): void
 
     public function getAvailableTokens(float $now): int
     {
-        $elapsed = $now - $this->timer;
+        $elapsed = max(0, $now - $this->timer);
 
         return min($this->burstSize, $this->tokens + $this->rate->calculateNewTokensDuringInterval($elapsed));
     }

diff --git a/Policy/TokenBucketLimiter.php b/Policy/TokenBucketLimiter.php
@@ -88,10 +88,10 @@ 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(0);
-                $bucket->setTimer($now + $waitDuration);
+                $bucket->setTokens($availableTokens - $tokens);
+                $bucket->setTimer($now);
 
-                $reservation = new Reservation($bucket->getTimer(), new RateLimit(0, \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false, $this->maxBurst));
+                $reservation = new Reservation($now + $waitDuration, new RateLimit(0, \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false, $this->maxBurst));
             }
 
             $this->storage->save($bucket);

diff --git a/Policy/Window.php b/Policy/Window.php
@@ -68,11 +68,6 @@ public function getHitCount(): int
 
     public function getAvailableTokens(float $now)
     {
-        // if timer is in future, there are no tokens available anymore
-        if ($this->timer > $now) {
-            return 0;
-        }
-
         // if now is more than the window interval in the past, all tokens are available
         if (($now - $this->timer) > $this->intervalInSeconds) {
             return $this->maxSize;

diff --git a/Tests/Policy/FixedWindowLimiterTest.php b/Tests/Policy/FixedWindowLimiterTest.php
@@ -14,6 +14,7 @@
 use PHPUnit\Framework\TestCase;
 use Symfony\Bridge\PhpUnit\ClockMock;
 use Symfony\Component\RateLimiter\Policy\FixedWindowLimiter;
+use Symfony\Component\RateLimiter\Policy\Window;
 use Symfony\Component\RateLimiter\RateLimit;
 use Symfony\Component\RateLimiter\Storage\InMemoryStorage;
 use Symfony\Component\RateLimiter\Tests\Resources\DummyWindow;
@@ -90,6 +91,19 @@ public function testWrongWindowFromCache()
         $this->assertEquals(9, $rateLimit->getRemainingTokens());
     }
 
+    public function testWindowResilientToTimeShifting()
+    {
+        $serverOneClock = microtime(true) - 1;
+        $serverTwoClock = microtime(true) + 1;
+        $window = new Window('id', 300, 100, $serverTwoClock);
+        $this->assertSame(100, $window->getAvailableTokens($serverTwoClock));
+        $this->assertSame(100, $window->getAvailableTokens($serverOneClock));
+
+        $window = new Window('id', 300, 100, $serverOneClock);
+        $this->assertSame(100, $window->getAvailableTokens($serverTwoClock));
+        $this->assertSame(100, $window->getAvailableTokens($serverOneClock));
+    }
+
     private function createLimiter(): FixedWindowLimiter
     {
         return new FixedWindowLimiter('test', 10, new \DateInterval('PT1M'), $this->storage);

diff --git a/Tests/Policy/TokenBucketLimiterTest.php b/Tests/Policy/TokenBucketLimiterTest.php
@@ -114,6 +114,20 @@ public function testWrongWindowFromCache()
         $this->assertEquals(9, $rateLimit->getRemainingTokens());
     }
 
+    public function testBucketResilientToTimeShifting()
+    {
+        $serverOneClock = microtime(true) - 1;
+        $serverTwoClock = microtime(true) + 1;
+
+        $bucket = new TokenBucket('id', 100, new Rate(\DateInterval::createFromDateString('5 minutes'), 10), $serverTwoClock);
+        $this->assertSame(100, $bucket->getAvailableTokens($serverTwoClock));
+        $this->assertSame(100, $bucket->getAvailableTokens($serverOneClock));
+
+        $bucket = new TokenBucket('id', 100, new Rate(\DateInterval::createFromDateString('5 minutes'), 10), $serverOneClock);
+        $this->assertSame(100, $bucket->getAvailableTokens($serverTwoClock));
+        $this->assertSame(100, $bucket->getAvailableTokens($serverOneClock));
+    }
+
     private function createLimiter($initialTokens = 10, Rate $rate = null)
     {
         return new TokenBucketLimiter('test', $initialTokens, $rate ?? Rate::perSecond(10), $this->storage);