calculateTimeForTokens for SlidingWindow

Jeroeny · Jeroeny · commit ca7ff0aa86f2 · 2022-11-22T14:23:56.000+01:00
diff --git a/src/Symfony/Component/RateLimiter/Policy/SlidingWindow.php b/src/Symfony/Component/RateLimiter/Policy/SlidingWindow.php
@@ -89,6 +89,27 @@ public function getRetryAfter(): \DateTimeImmutable
         return \DateTimeImmutable::createFromFormat('U.u', sprintf('%.6F', $this->windowEndAt));
     }
 
+    public function calculateTimeForTokens(int $maxSize, int $tokens): int
+    {
+        $remaining = $maxSize - $this->getHitCount();
+        if ($remaining >= $tokens) {
+            return 0;
+        }
+
+        $startOfWindow = $this->windowEndAt - $this->intervalInSeconds;
+        $percentOfCurrentTimeFrame = min((microtime(true) - $startOfWindow) / $this->intervalInSeconds, 1);
+        $releasable = $maxSize - floor($this->hitCountForLastWindow * (1 - $percentOfCurrentTimeFrame));
+        $remainingWindow = (microtime(true) - $startOfWindow) - $this->intervalInSeconds;
+        $timePerToken = $remainingWindow / $releasable;
+        $needed = $tokens - $remaining;
+
+        if ($releasable <= $needed) {
+            return (int) ceil($needed * $timePerToken);
+        }
+
+        return (int) ($this->windowEndAt - microtime(true)) + ceil(($needed - $releasable) * ($this->intervalInSeconds / $maxSize));
+    }
+
     public function __serialize(): array
     {
         return [
diff --git a/src/Symfony/Component/RateLimiter/Policy/SlidingWindowLimiter.php b/src/Symfony/Component/RateLimiter/Policy/SlidingWindowLimiter.php
@@ -12,8 +12,8 @@
 namespace Symfony\Component\RateLimiter\Policy;
 
 use Symfony\Component\Lock\LockInterface;
-use Symfony\Component\RateLimiter\Exception\ReserveNotSupportedException;
 use Symfony\Component\RateLimiter\LimiterInterface;
+use Symfony\Component\RateLimiter\Exception\MaxWaitDurationExceededException;
 use Symfony\Component\RateLimiter\RateLimit;
 use Symfony\Component\RateLimiter\Reservation;
 use Symfony\Component\RateLimiter\Storage\StorageInterface;
@@ -48,11 +48,10 @@ public function __construct(string $id, int $limit, \DateInterval $interval, Sto
 
     public function reserve(int $tokens = 1, float $maxTime = null): Reservation
     {
-        throw new ReserveNotSupportedException(__CLASS__);
-    }
+        if ($tokens > $this->limit) {
+            throw new \InvalidArgumentException(sprintf('Cannot reserve more tokens (%d) than the size of the rate limiter (%d).', $tokens, $this->limit));
+        }
 
-    public function consume(int $tokens = 1): RateLimit
-    {
         $this->lock?->acquire(true);
 
         try {
@@ -63,22 +62,43 @@ public function consume(int $tokens = 1): RateLimit
                 $window = SlidingWindow::createFromPreviousWindow($window, $this->interval);
             }
 
+            $now = microtime(true);
             $hitCount = $window->getHitCount();
             $availableTokens = $this->getAvailableTokens($hitCount);
-            if ($availableTokens < $tokens) {
-                return new RateLimit($availableTokens, $window->getRetryAfter(), false, $this->limit);
-            }
+            if ($availableTokens >= $tokens) {
+                $window->add($tokens);
 
-            $window->add($tokens);
+                $reservation = new Reservation($now, new RateLimit($this->getAvailableTokens($window->getHitCount()), \DateTimeImmutable::createFromFormat('U', floor($now)), true, $this->limit));
+            } else {
+                $waitDuration = $window->calculateTimeForTokens($this->limit, max(1, $tokens));
+
+                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 RateLimit($this->getAvailableTokens($window->getHitCount()), \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false, $this->limit));
+                }
+
+                $window->add($tokens);
+
+                $reservation = new Reservation($now + $waitDuration, new RateLimit($this->getAvailableTokens($window->getHitCount()), \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false, $this->limit));
+            }
 
             if (0 < $tokens) {
                 $this->storage->save($window);
             }
-
-            return new RateLimit($this->getAvailableTokens($window->getHitCount()), $window->getRetryAfter(), true, $this->limit);
         } finally {
             $this->lock?->release();
         }
+
+        return $reservation;
+    }
+
+    public function consume(int $tokens = 1): RateLimit
+    {
+        try {
+            return $this->reserve($tokens, 0)->getRateLimit();
+        } catch (MaxWaitDurationExceededException $e) {
+            return $e->getRateLimit();
+        }
     }
 
     private function getAvailableTokens(int $hitCount): int
diff --git a/src/Symfony/Component/RateLimiter/Tests/Policy/SlidingWindowLimiterTest.php b/src/Symfony/Component/RateLimiter/Tests/Policy/SlidingWindowLimiterTest.php
@@ -13,7 +13,6 @@
 
 use PHPUnit\Framework\TestCase;
 use Symfony\Bridge\PhpUnit\ClockMock;
-use Symfony\Component\RateLimiter\Exception\ReserveNotSupportedException;
 use Symfony\Component\RateLimiter\Policy\SlidingWindowLimiter;
 use Symfony\Component\RateLimiter\RateLimit;
 use Symfony\Component\RateLimiter\Storage\InMemoryStorage;
@@ -66,14 +65,17 @@ public function testWaitIntervalOnConsumeOverLimit()
 
         $start = microtime(true);
         $rateLimit->wait(); // wait 12 seconds
-        $this->assertEqualsWithDelta($start + 12, microtime(true), 1);
+        $this->assertEqualsWithDelta($start + (12 / 5), microtime(true), 1);
+        $this->assertTrue($limiter->consume()->isAccepted());
     }
 
     public function testReserve()
     {
-        $this->expectException(ReserveNotSupportedException::class);
+        $limiter = $this->createLimiter();
+        $limiter->consume(8);
 
-        $this->createLimiter()->reserve();
+        // 2 over the limit, causing the WaitDuration to become 2/10th of the 12s interval
+        $this->assertEqualsWithDelta(12 / 5, $limiter->reserve(4)->getWaitDuration(), 1);
     }
 
     public function testPeekConsume()
