calculateTimeForTokens for SlidingWindow

Jeroeny · Jeroeny · commit 459b33227cf5 · 2022-09-12T16:39:10.000+02: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 ceil($needed * $timePerToken);
+        }
+
+        return ($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
@@ -13,7 +13,7 @@
 
 use Symfony\Component\Lock\LockInterface;
 use Symfony\Component\Lock\NoLock;
-use Symfony\Component\RateLimiter\Exception\ReserveNotSupportedException;
+use Symfony\Component\RateLimiter\Exception\MaxWaitDurationExceededException;
 use Symfony\Component\RateLimiter\LimiterInterface;
 use Symfony\Component\RateLimiter\RateLimit;
 use Symfony\Component\RateLimiter\Reservation;
@@ -49,11 +49,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 {
@@ -64,22 +63,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
