[RateLimit] Test and fix peeking behavior on rate limit policies

wouterj · Dec 28, 2023 · e4a8c33 · e4a8c33
1 parent 843f7bc
commit e4a8c33
Show file tree

Hide file tree

Showing 6 changed files with 66 additions and 8 deletions.
diff --git a/src/Symfony/Component/RateLimiter/Policy/FixedWindowLimiter.php b/src/Symfony/Component/RateLimiter/Policy/FixedWindowLimiter.php
@@ -59,12 +59,15 @@ public function reserve(int $tokens = 1, float $maxTime = null): Reservation
             $now = microtime(true);
             $availableTokens = $window->getAvailableTokens($now);
 
-            if ($availableTokens >= max(1, $tokens)) {
+            if (0 === $tokens) {
+                $waitDuration = $window->calculateTimeForTokens(1, $now);
+                $reservation = new Reservation($now + $waitDuration, new RateLimit($window->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), true, $this->limit));
+            } elseif ($availableTokens >= $tokens) {
                 $window->add($tokens, $now);
 
                 $reservation = new Reservation($now, new RateLimit($window->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now)), true, $this->limit));
             } else {
-                $waitDuration = $window->calculateTimeForTokens(max(1, $tokens), $now);
+                $waitDuration = $window->calculateTimeForTokens($tokens, $now);
 
                 if (null !== $maxTime && $waitDuration > $maxTime) {
                     // process needs to wait longer than set interval

diff --git a/src/Symfony/Component/RateLimiter/Policy/SlidingWindowLimiter.php b/src/Symfony/Component/RateLimiter/Policy/SlidingWindowLimiter.php
@@ -69,12 +69,13 @@ public function consume(int $tokens = 1): RateLimit
                 return new RateLimit($availableTokens, $window->getRetryAfter(), false, $this->limit);
             }
 
-            $window->add($tokens);
-
-            if (0 < $tokens) {
-                $this->storage->save($window);
+            if (0 === $tokens) {
+                return new RateLimit($availableTokens, $availableTokens ? \DateTimeImmutable::createFromFormat('U.u', sprintf('%.6F', microtime(true))) : $window->getRetryAfter(), true, $this->limit);
             }
 
+            $window->add($tokens);
+            $this->storage->save($window);
+
             return new RateLimit($this->getAvailableTokens($window->getHitCount()), $window->getRetryAfter(), true, $this->limit);
         } finally {
             $this->lock?->release();

diff --git a/src/Symfony/Component/RateLimiter/Policy/TokenBucketLimiter.php b/src/Symfony/Component/RateLimiter/Policy/TokenBucketLimiter.php
@@ -67,11 +67,20 @@ public function reserve(int $tokens = 1, float $maxTime = null): Reservation
             $now = microtime(true);
             $availableTokens = $bucket->getAvailableTokens($now);
 
-            if ($availableTokens >= max(1, $tokens)) {
+            if ($availableTokens >= $tokens) {
                 // tokens are now available, update bucket
                 $bucket->setTokens($availableTokens - $tokens);
 
-                $reservation = new Reservation($now, new RateLimit($bucket->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now)), true, $this->maxBurst));
+                if (0 === $availableTokens) {
+                    // This means 0 tokens where consumed (discouraged in most cases).
+                    // Return the first time a new token is available
+                    $waitDuration = $this->rate->calculateTimeForTokens(1);
+                    $waitTime = \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration));
+                } else {
+                    $waitTime = \DateTimeImmutable::createFromFormat('U', floor($now));
+                }
+
+                $reservation = new Reservation($now, new RateLimit($bucket->getAvailableTokens($now), $waitTime, true, $this->maxBurst));
             } else {
                 $remainingTokens = $tokens - $availableTokens;
                 $waitDuration = $this->rate->calculateTimeForTokens($remainingTokens);

diff --git a/src/Symfony/Component/RateLimiter/Tests/Policy/FixedWindowLimiterTest.php b/src/Symfony/Component/RateLimiter/Tests/Policy/FixedWindowLimiterTest.php
@@ -123,7 +123,21 @@ public function testPeekConsume()
             $rateLimit = $limiter->consume(0);
             $this->assertSame(10, $rateLimit->getLimit());
             $this->assertTrue($rateLimit->isAccepted());
+            $this->assertEquals(
+                \DateTimeImmutable::createFromFormat('U', (string) floor(microtime(true))),
+                $rateLimit->getRetryAfter()
+            );
         }
+
+        $limiter->consume();
+
+        $rateLimit = $limiter->consume(0);
+        $this->assertEquals(0, $rateLimit->getRemainingTokens());
+        $this->assertTrue($rateLimit->isAccepted());
+        $this->assertEquals(
+            \DateTimeImmutable::createFromFormat('U', (string) floor(microtime(true) + 60)),
+            $rateLimit->getRetryAfter()
+        );
     }
 
     public static function provideConsumeOutsideInterval(): \Generator

diff --git a/src/Symfony/Component/RateLimiter/Tests/Policy/SlidingWindowLimiterTest.php b/src/Symfony/Component/RateLimiter/Tests/Policy/SlidingWindowLimiterTest.php
@@ -31,6 +31,7 @@ protected function setUp(): void
 
         ClockMock::register(InMemoryStorage::class);
         ClockMock::register(RateLimit::class);
+        ClockMock::register(SlidingWindowLimiter::class);
     }
 
     public function testConsume()
@@ -82,11 +83,26 @@ public function testPeekConsume()
 
         $limiter->consume(9);
 
+        // peek by consuming 0 tokens twice (making sure peeking doesn't claim a token)
         for ($i = 0; $i < 2; ++$i) {
             $rateLimit = $limiter->consume(0);
             $this->assertTrue($rateLimit->isAccepted());
             $this->assertSame(10, $rateLimit->getLimit());
+            $this->assertEquals(
+                \DateTimeImmutable::createFromFormat('U.u', sprintf('%.6F', microtime(true))),
+                $rateLimit->getRetryAfter()
+            );
         }
+
+        $limiter->consume();
+
+        $rateLimit = $limiter->consume(0);
+        $this->assertEquals(0, $rateLimit->getRemainingTokens());
+        $this->assertTrue($rateLimit->isAccepted());
+        $this->assertEquals(
+            \DateTimeImmutable::createFromFormat('U.u', sprintf('%.6F', microtime(true) + 12)),
+            $rateLimit->getRetryAfter()
+        );
     }
 
     private function createLimiter(): SlidingWindowLimiter

diff --git a/src/Symfony/Component/RateLimiter/Tests/Policy/TokenBucketLimiterTest.php b/src/Symfony/Component/RateLimiter/Tests/Policy/TokenBucketLimiterTest.php
@@ -134,11 +134,26 @@ public function testPeekConsume()
 
         $limiter->consume(9);
 
+        // peek by consuming 0 tokens twice (making sure peeking doesn't claim a token)
         for ($i = 0; $i < 2; ++$i) {
             $rateLimit = $limiter->consume(0);
             $this->assertTrue($rateLimit->isAccepted());
             $this->assertSame(10, $rateLimit->getLimit());
+            $this->assertEquals(
+                \DateTimeImmutable::createFromFormat('U', (string) floor(microtime(true))),
+                $rateLimit->getRetryAfter()
+            );
         }
+
+        $limiter->consume();
+
+        $rateLimit = $limiter->consume(0);
+        $this->assertEquals(0, $rateLimit->getRemainingTokens());
+        $this->assertTrue($rateLimit->isAccepted());
+        $this->assertEquals(
+            \DateTimeImmutable::createFromFormat('U', (string) floor(microtime(true) + 1)),
+            $rateLimit->getRetryAfter()
+        );
     }
 
     public function testBucketRefilledWithStrictFrequency()