Add reserve() to the LimiterInterfaces

See symfony/symfony-docs#14370 (comment)

Author: wouterj

src/Symfony/Component/RateLimiter/CompoundLimiter.php

@@ -11,6 +11,8 @@
 
 namespace Symfony\Component\RateLimiter;
 
+use Symfony\Component\RateLimiter\Exception\MaxWaitDurationExceededException;
+
 /**
  * @author Wouter de Jong <[email protected]>
  *
@@ -31,6 +33,11 @@ public function __construct(array $limiters)
         $this->limiters = $limiters;
     }
 
+    public function reserve(int $tokens = 1, ?float $maxTime = null): Reservation
+    {
+        // TODO: Implement reserve() method.
+    }
+
     public function consume(int $tokens = 1): Limit
     {
         $minimalLimit = null;

src/Symfony/Component/RateLimiter/Exception/MaxWaitDurationExceededException.php

@@ -11,11 +11,27 @@
 
 namespace Symfony\Component\RateLimiter\Exception;
 
+use Symfony\Component\RateLimiter\Limit;
+use Throwable;
+
 /**
  * @author Wouter de Jong <[email protected]>
  *
  * @experimental in 5.2
  */
 class MaxWaitDurationExceededException extends \RuntimeException
 {
+    private $limit;
+
+    public function __construct(string $message, Limit $limit, int $code = 0, Throwable $previous = null)
+    {
+        parent::__construct($message, $code, $previous);
+
+        $this->limit = $limit;
+    }
+
+    public function getLimit(): Limit
+    {
+        return $this->limit;
+    }
 }

src/Symfony/Component/RateLimiter/FixedWindowLimiter.php

@@ -13,6 +13,7 @@
 
 use Symfony\Component\Lock\LockInterface;
 use Symfony\Component\Lock\NoLock;
+use Symfony\Component\RateLimiter\Exception\MaxWaitDurationExceededException;
 use Symfony\Component\RateLimiter\Storage\StorageInterface;
 use Symfony\Component\RateLimiter\Util\TimeUtil;
 
@@ -40,42 +41,64 @@ public function __construct(string $id, int $limit, \DateInterval $interval, Sto
         $this->interval = TimeUtil::dateIntervalToSeconds($interval);
     }
 
-    /**
-     * {@inheritdoc}
-     */
-    public function consume(int $tokens = 1): Limit
+    public function reserve(int $tokens = 1, ?float $maxTime = null): Reservation
     {
+        if ($tokens > $this->limit) {
+            throw new \InvalidArgumentException(sprintf('Cannot reserve more tokens (%d) than the size of the rate limiter (%d).', $tokens, $this->limit));
+        }
+
         $this->lock->acquire(true);
 
         try {
             $window = $this->storage->fetch($this->id);
             if (!$window instanceof Window) {
-                $window = new Window($this->id, $this->interval);
+                $window = new Window($this->id, $this->interval, $this->limit);
             }
 
-            $hitCount = $window->getHitCount();
-            $availableTokens = $this->getAvailableTokens($hitCount);
-            $windowStart = \DateTimeImmutable::createFromFormat('U', time());
-            if ($availableTokens < $tokens) {
-                return new Limit($availableTokens, $this->getRetryAfter($windowStart), false);
-            }
+            $now = microtime(true);
+            $availableTokens = $window->getAvailableTokens($now);
+            if ($availableTokens >= $tokens) {
+                $window->add($tokens);
 
-            $window->add($tokens);
-            $this->storage->save($window);
+                $reservation = new Reservation($now, new Limit($window->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now)), true));
+            } else {
+                $remainingTokens = $tokens - $availableTokens;
+                $waitDuration = $window->calculateTimeForTokens($remainingTokens);
+
+                if (null !== $maxTime && $waitDuration > $maxTime) {
+                    // process needs to wait longer than set interval
+                    throw new MaxWaitDurationExceededException(
+                        sprintf('The rate limiter wait time ("%d" seconds) is longer than the provided maximum time ("%d" seconds).', $waitDuration, $maxTime),
+                        new Limit($window->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false)
+                    );
+                }
 
-            return new Limit($this->getAvailableTokens($window->getHitCount()), $this->getRetryAfter($windowStart), true);
+                $window->add($tokens);
+
+                $reservation = new Reservation($now + $waitDuration, new Limit($window->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false));
+            }
+            $this->storage->save($window);
         } finally {
             $this->lock->release();
         }
+
+        return $reservation;
     }
 
-    public function getAvailableTokens(int $hitCount): int
+    /**
+     * {@inheritdoc}
+     */
+    public function consume(int $tokens = 1): Limit
     {
-        return $this->limit - $hitCount;
+        try {
+            return $this->reserve($tokens, 0)->getLimit();
+        } catch (MaxWaitDurationExceededException $e) {
+            return $e->getLimit();
+        }
     }
 
-    private function getRetryAfter(\DateTimeImmutable $windowStart): \DateTimeImmutable
+    public function getAvailableTokens(int $hitCount): int
     {
-        return $windowStart->add(new \DateInterval(sprintf('PT%sS', $this->interval)));
+        return $this->limit - $hitCount;
     }
 }

src/Symfony/Component/RateLimiter/LimiterInterface.php

@@ -11,13 +11,30 @@
 
 namespace Symfony\Component\RateLimiter;
 
+use Symfony\Component\RateLimiter\Exception\MaxWaitDurationExceededException;
+
 /**
  * @author Wouter de Jong <[email protected]>
  *
  * @experimental in 5.2
  */
 interface LimiterInterface
 {
+    /**
+     * Waits until the required number of tokens is available.
+     *
+     * The reserved tokens will be taken into account when calculating
+     * future token consumptions. Do not use this method if you intend
+     * to skip this process.
+     *
+     * @param int   $tokens  the number of tokens required
+     * @param float $maxTime maximum accepted waiting time in seconds
+     *
+     * @throws MaxWaitDurationExceededException if $maxTime is set and the process needs to wait longer than its value (in seconds)
+     * @throws \InvalidArgumentException        if $tokens is larger than the maximum burst size
+     */
+    public function reserve(int $tokens = 1, ?float $maxTime = null): Reservation;
+
     /**
      * Use this method if you intend to drop if the required number
      * of tokens is unavailable.

src/Symfony/Component/RateLimiter/Reservation.php

@@ -19,13 +19,15 @@
 final class Reservation
 {
     private $timeToAct;
+    private $limit;
 
     /**
      * @param float $timeToAct Unix timestamp in seconds when this reservation should act
      */
-    public function __construct(float $timeToAct)
+    public function __construct(float $timeToAct, Limit $limit)
     {
         $this->timeToAct = $timeToAct;
+        $this->limit = $limit;
     }
 
     public function getTimeToAct(): float
@@ -38,6 +40,11 @@ public function getWaitDuration(): float
         return max(0, (-microtime(true)) + $this->timeToAct);
     }
 
+    public function getLimit(): Limit
+    {
+        return $this->limit;
+    }
+
     public function wait(): void
     {
         usleep($this->getWaitDuration() * 1e6);

src/Symfony/Component/RateLimiter/Tests/CompoundLimiterTest.php

@@ -38,17 +38,14 @@ public function testConsume()
         $limiter3 = $this->createLimiter(12, new \DateInterval('PT30S'));
         $limiter = new CompoundLimiter([$limiter1, $limiter2, $limiter3]);
 
-        // Reach limiter 1 limit, verify that limiter2 available tokens reduced by 5 and fetch successfully limiter 1
-        $this->assertEquals(3, $limiter->consume(5)->getRemainingTokens(), 'Limiter 1 reached the limit');
+        $this->assertEquals(0, $limiter->consume(4)->getRemainingTokens(), 'Limiter 1 reached the limit');
         sleep(1); // reset limiter1's window
-        $this->assertTrue($limiter->consume(2)->isAccepted());
+        $this->assertTrue($limiter->consume(3)->isAccepted());
 
-        // Reach limiter 2 limit, verify that limiter2 available tokens reduced by 5 and and fetch successfully
         $this->assertEquals(0, $limiter->consume()->getRemainingTokens(), 'Limiter 2 has no remaining tokens left');
-        sleep(9); // reset limiter2's window
+        sleep(10); // reset limiter2's window
         $this->assertTrue($limiter->consume(3)->isAccepted());
 
-        // Reach limiter 3 limit, verify that limiter2 available tokens reduced by 5 and fetch successfully
         $this->assertEquals(0, $limiter->consume()->getRemainingTokens(), 'Limiter 3 reached the limit');
         sleep(20); // reset limiter3's window
         $this->assertTrue($limiter->consume()->isAccepted());

src/Symfony/Component/RateLimiter/Tests/FixedWindowLimiterTest.php

@@ -60,7 +60,7 @@ public function testConsumeOutsideInterval()
         sleep(10);
         $limit = $limiter->consume(10);
         $this->assertEquals(0, $limit->getRemainingTokens());
-        $this->assertEquals(time() + 60, $limit->getRetryAfter()->getTimestamp());
+        $this->assertTrue($limit->isAccepted());
     }
 
     public function testWrongWindowFromCache()

src/Symfony/Component/RateLimiter/Tests/Storage/CacheStorageTest.php

@@ -31,23 +31,22 @@ protected function setUp(): void
     public function testSave()
     {
         $cacheItem = $this->createMock(CacheItemInterface::class);
-        $cacheItem->expects($this->once())->method('expiresAfter')->with(10);
+        $cacheItem->expects($this->exactly(2))->method('expiresAfter')->with(10);
 
         $this->pool->expects($this->any())->method('getItem')->with(sha1('test'))->willReturn($cacheItem);
         $this->pool->expects($this->exactly(2))->method('save')->with($cacheItem);
 
-        $window = new Window('test', 10);
+        $window = new Window('test', 10, 20);
         $this->storage->save($window);
 
-        // test that expiresAfter is only called when getExpirationAt() does not return null
         $window = unserialize(serialize($window));
         $this->storage->save($window);
     }
 
     public function testFetchExistingState()
     {
         $cacheItem = $this->createMock(CacheItemInterface::class);
-        $window = new Window('test', 10);
+        $window = new Window('test', 10, 20);
         $cacheItem->expects($this->any())->method('get')->willReturn($window);
         $cacheItem->expects($this->any())->method('isHit')->willReturn(true);
 

src/Symfony/Component/RateLimiter/TokenBucket.php

@@ -14,6 +14,7 @@
 /**
  * @author Wouter de Jong <[email protected]>
  *
+ * @internal
  * @experimental in 5.2
  */
 final class TokenBucket implements LimiterStateInterface

src/Symfony/Component/RateLimiter/TokenBucketLimiter.php

@@ -74,22 +74,24 @@ public function reserve(int $tokens = 1, ?float $maxTime = null): Reservation
                 $bucket->setTokens($availableTokens - $tokens);
                 $bucket->setTimer($now);
 
-                $reservation = new Reservation($now);
+                $reservation = new Reservation($now, new Limit($bucket->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now)), true));
             } else {
                 $remainingTokens = $tokens - $availableTokens;
                 $waitDuration = $this->rate->calculateTimeForTokens($remainingTokens);
 
                 if (null !== $maxTime && $waitDuration > $maxTime) {
                     // process needs to wait longer than set interval
-                    throw new MaxWaitDurationExceededException(sprintf('The rate limiter wait time ("%d" seconds) is longer than the provided maximum time ("%d" seconds).', $waitDuration, $maxTime));
+                    $limit = new Limit($availableTokens, \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false);
+
+                    throw new MaxWaitDurationExceededException(sprintf('The rate limiter wait time ("%d" seconds) is longer than the provided maximum time ("%d" seconds).', $waitDuration, $maxTime), $limit);
                 }
 
                 // 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);
 
-                $reservation = new Reservation($bucket->getTimer());
+                $reservation = new Reservation($bucket->getTimer(), new Limit(0, \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false));
             }
 
             $this->storage->save($bucket);
@@ -105,18 +107,10 @@ public function reserve(int $tokens = 1, ?float $maxTime = null): Reservation
      */
     public function consume(int $tokens = 1): Limit
     {
-        $bucket = $this->storage->fetch($this->id);
-        if (!$bucket instanceof TokenBucket) {
-            $bucket = new TokenBucket($this->id, $this->maxBurst, $this->rate);
-        }
-        $now = microtime(true);
-
         try {
-            $this->reserve($tokens, 0);
-
-            return new Limit($bucket->getAvailableTokens($now) - $tokens, $this->rate->calculateNextTokenAvailability(), true);
+            return $this->reserve($tokens, 0)->getLimit();
         } catch (MaxWaitDurationExceededException $e) {
-            return new Limit($bucket->getAvailableTokens($now), $this->rate->calculateNextTokenAvailability(), false);
+            return $e->getLimit();
         }
     }
 }