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\ReserveNotSupportedException;
+
 /**
  * @author Wouter de Jong <wouter@wouterj.nl>
  *
@@ -31,6 +33,11 @@ public function __construct(array $limiters)
         $this->limiters = $limiters;
     }
 
+    public function reserve(int $tokens = 1, ?float $maxTime = null): Reservation
+    {
+        throw new ReserveNotSupportedException(__CLASS__);
+    }
+
     public function consume(int $tokens = 1): Limit
     {
         $minimalLimit = null;

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

@@ -11,11 +11,26 @@
 
 namespace Symfony\Component\RateLimiter\Exception;
 
+use Symfony\Component\RateLimiter\Limit;
+
 /**
  * @author Wouter de Jong <wouter@wouterj.nl>
  *
  * @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/Exception/ReserveNotSupportedException.php

@@ -0,0 +1,25 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <fabien@symfony.com>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\Component\RateLimiter\Exception;
+
+/**
+ * @author Wouter de Jong <wouter@wouterj.nl>
+ *
+ * @experimental in 5.2
+ */
+class ReserveNotSupportedException extends \BadMethodCallException
+{
+    public function __construct(string $limiterClass, int $code = 0, ?\Throwable $previous = null)
+    {
+        parent::__construct(sprintf('Reserving tokens is not supported by "%s".', $limiterClass), $code, $previous);
+    }
+}

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,61 @@ 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,32 @@
 
 namespace Symfony\Component\RateLimiter;
 
+use Symfony\Component\RateLimiter\Exception\MaxWaitDurationExceededException;
+use Symfony\Component\RateLimiter\Exception\ReserveNotSupportedException;
+
 /**
  * @author Wouter de Jong <wouter@wouterj.nl>
  *
  * @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 ReserveNotSupportedException     if this limiter implementation doesn't support reserving tokens
+     * @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/SlidingWindowLimiter.php

@@ -13,6 +13,7 @@
 
 use Symfony\Component\Lock\LockInterface;
 use Symfony\Component\Lock\NoLock;
+use Symfony\Component\RateLimiter\Exception\ReserveNotSupportedException;
 use Symfony\Component\RateLimiter\Storage\StorageInterface;
 use Symfony\Component\RateLimiter\Util\TimeUtil;
 
@@ -67,6 +68,11 @@ public function __construct(string $id, int $limit, \DateInterval $interval, Sto
         $this->interval = TimeUtil::dateIntervalToSeconds($interval);
     }
 
+    public function reserve(int $tokens = 1, ?float $maxTime = null): Reservation
+    {
+        throw new ReserveNotSupportedException(__CLASS__);
+    }
+
     /**
      * {@inheritdoc}
      */

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

@@ -14,6 +14,7 @@
 use PHPUnit\Framework\TestCase;
 use Symfony\Bridge\PhpUnit\ClockMock;
 use Symfony\Component\RateLimiter\CompoundLimiter;
+use Symfony\Component\RateLimiter\Exception\ReserveNotSupportedException;
 use Symfony\Component\RateLimiter\FixedWindowLimiter;
 use Symfony\Component\RateLimiter\Storage\InMemoryStorage;
 
@@ -38,22 +39,26 @@ 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());
     }
 
+    public function testReserve()
+    {
+        $this->expectException(ReserveNotSupportedException::class);
+
+        (new CompoundLimiter([$this->createLimiter(4, new \DateInterval('PT1S'))]))->reserve();
+    }
+
     private function createLimiter(int $limit, \DateInterval $interval): FixedWindowLimiter
     {
         return new FixedWindowLimiter('test'.$limit, $limit, $interval, $this->storage);

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/SlidingWindowLimiterTest.php

@@ -13,6 +13,7 @@
 
 use PHPUnit\Framework\TestCase;
 use Symfony\Bridge\PhpUnit\ClockMock;
+use Symfony\Component\RateLimiter\Exception\ReserveNotSupportedException;
 use Symfony\Component\RateLimiter\SlidingWindowLimiter;
 use Symfony\Component\RateLimiter\Storage\InMemoryStorage;
 
@@ -50,6 +51,13 @@ public function testConsume()
         $this->assertTrue($limit->isAccepted());
     }
 
+    public function testReserve()
+    {
+        $this->expectException(ReserveNotSupportedException::class);
+
+        $this->createLimiter()->reserve();
+    }
+
     private function createLimiter(): SlidingWindowLimiter
     {
         return new SlidingWindowLimiter('test', 10, new \DateInterval('PT12S'), $this->storage);