|
13 | 13 |
|
14 | 14 | use Symfony\Component\Lock\LockInterface;
|
15 | 15 | use Symfony\Component\Lock\NoLock;
|
| 16 | +use Symfony\Component\RateLimiter\Exception\MaxWaitDurationExceededException; |
16 | 17 | use Symfony\Component\RateLimiter\Storage\StorageInterface;
|
17 | 18 | use Symfony\Component\RateLimiter\Util\TimeUtil;
|
18 | 19 |
|
@@ -40,42 +41,61 @@ public function __construct(string $id, int $limit, \DateInterval $interval, Sto
|
40 | 41 | $this->interval = TimeUtil::dateIntervalToSeconds($interval);
|
41 | 42 | }
|
42 | 43 |
|
43 |
| - /** |
44 |
| - * {@inheritdoc} |
45 |
| - */ |
46 |
| - public function consume(int $tokens = 1): Limit |
| 44 | + public function reserve(int $tokens = 1, ?float $maxTime = null): Reservation |
47 | 45 | {
|
| 46 | + if ($tokens > $this->limit) { |
| 47 | + throw new \InvalidArgumentException(sprintf('Cannot reserve more tokens (%d) than the size of the rate limiter (%d).', $tokens, $this->limit)); |
| 48 | + } |
| 49 | + |
48 | 50 | $this->lock->acquire(true);
|
49 | 51 |
|
50 | 52 | try {
|
51 | 53 | $window = $this->storage->fetch($this->id);
|
52 | 54 | if (!$window instanceof Window) {
|
53 |
| - $window = new Window($this->id, $this->interval); |
| 55 | + $window = new Window($this->id, $this->interval, $this->limit); |
54 | 56 | }
|
55 | 57 |
|
56 |
| - $hitCount = $window->getHitCount(); |
57 |
| - $availableTokens = $this->getAvailableTokens($hitCount); |
58 |
| - $windowStart = \DateTimeImmutable::createFromFormat('U', time()); |
59 |
| - if ($availableTokens < $tokens) { |
60 |
| - return new Limit($availableTokens, $this->getRetryAfter($windowStart), false); |
61 |
| - } |
| 58 | + $now = microtime(true); |
| 59 | + $availableTokens = $window->getAvailableTokens($now); |
| 60 | + if ($availableTokens >= $tokens) { |
| 61 | + $window->add($tokens); |
62 | 62 |
|
63 |
| - $window->add($tokens); |
64 |
| - $this->storage->save($window); |
| 63 | + $reservation = new Reservation($now, new Limit($window->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now)), true)); |
| 64 | + } else { |
| 65 | + $remainingTokens = $tokens - $availableTokens; |
| 66 | + $waitDuration = $window->calculateTimeForTokens($remainingTokens); |
| 67 | + |
| 68 | + if (null !== $maxTime && $waitDuration > $maxTime) { |
| 69 | + // process needs to wait longer than set interval |
| 70 | + 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)); |
| 71 | + } |
65 | 72 |
|
66 |
| - return new Limit($this->getAvailableTokens($window->getHitCount()), $this->getRetryAfter($windowStart), true); |
| 73 | + $window->add($tokens); |
| 74 | + |
| 75 | + $reservation = new Reservation($now + $waitDuration, new Limit($window->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false)); |
| 76 | + } |
| 77 | + $this->storage->save($window); |
67 | 78 | } finally {
|
68 | 79 | $this->lock->release();
|
69 | 80 | }
|
| 81 | + |
| 82 | + return $reservation; |
70 | 83 | }
|
71 | 84 |
|
72 |
| - public function getAvailableTokens(int $hitCount): int |
| 85 | + /** |
| 86 | + * {@inheritdoc} |
| 87 | + */ |
| 88 | + public function consume(int $tokens = 1): Limit |
73 | 89 | {
|
74 |
| - return $this->limit - $hitCount; |
| 90 | + try { |
| 91 | + return $this->reserve($tokens, 0)->getLimit(); |
| 92 | + } catch (MaxWaitDurationExceededException $e) { |
| 93 | + return $e->getLimit(); |
| 94 | + } |
75 | 95 | }
|
76 | 96 |
|
77 |
| - private function getRetryAfter(\DateTimeImmutable $windowStart): \DateTimeImmutable |
| 97 | + public function getAvailableTokens(int $hitCount): int |
78 | 98 | {
|
79 |
| - return $windowStart->add(new \DateInterval(sprintf('PT%sS', $this->interval))); |
| 99 | + return $this->limit - $hitCount; |
80 | 100 | }
|
81 | 101 | }
|
0 commit comments