vendor/symfony/rate-limiter/Policy/TokenBucketLimiter.php line 56

Open in your IDE?
  1. <?php
  2. /*
  3. * This file is part of the Symfony package.
  4. *
  5. * (c) Fabien Potencier <fabien@symfony.com>
  6. *
  7. * For the full copyright and license information, please view the LICENSE
  8. * file that was distributed with this source code.
  9. */
  10. namespace Symfony\Component\RateLimiter\Policy;
  11. use Symfony\Component\Lock\LockInterface;
  12. use Symfony\Component\Lock\NoLock;
  13. use Symfony\Component\RateLimiter\Exception\MaxWaitDurationExceededException;
  14. use Symfony\Component\RateLimiter\LimiterInterface;
  15. use Symfony\Component\RateLimiter\RateLimit;
  16. use Symfony\Component\RateLimiter\Reservation;
  17. use Symfony\Component\RateLimiter\Storage\StorageInterface;
  18. /**
  19. * @author Wouter de Jong <wouter@wouterj.nl>
  20. *
  21. * @experimental in 5.3
  22. */
  23. final class TokenBucketLimiter implements LimiterInterface
  24. {
  25. use ResetLimiterTrait;
  26. private $maxBurst;
  27. private $rate;
  28. public function __construct(string $id, int $maxBurst, Rate $rate, StorageInterface $storage, LockInterface $lock = null)
  29. {
  30. $this->id = $id;
  31. $this->maxBurst = $maxBurst;
  32. $this->rate = $rate;
  33. $this->storage = $storage;
  34. $this->lock = $lock ?? new NoLock();
  35. }
  36. /**
  37. * Waits until the required number of tokens is available.
  38. *
  39. * The reserved tokens will be taken into account when calculating
  40. * future token consumptions. Do not use this method if you intend
  41. * to skip this process.
  42. *
  43. * @param int $tokens the number of tokens required
  44. * @param float $maxTime maximum accepted waiting time in seconds
  45. *
  46. * @throws MaxWaitDurationExceededException if $maxTime is set and the process needs to wait longer than its value (in seconds)
  47. * @throws \InvalidArgumentException if $tokens is larger than the maximum burst size
  48. */
  49. public function reserve(int $tokens = 1, float $maxTime = null): Reservation
  50. {
  51. if ($tokens > $this->maxBurst) {
  52. throw new \InvalidArgumentException(sprintf('Cannot reserve more tokens (%d) than the burst size of the rate limiter (%d).', $tokens, $this->maxBurst));
  53. }
  54. $this->lock->acquire(true);
  55. try {
  56. $bucket = $this->storage->fetch($this->id);
  57. if (!$bucket instanceof TokenBucket) {
  58. $bucket = new TokenBucket($this->id, $this->maxBurst, $this->rate);
  59. }
  60. $now = microtime(true);
  61. $availableTokens = $bucket->getAvailableTokens($now);
  62. if ($availableTokens >= $tokens) {
  63. // tokens are now available, update bucket
  64. $bucket->setTokens($availableTokens - $tokens);
  65. $bucket->setTimer($now);
  66. $reservation = new Reservation($now, new RateLimit($bucket->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now)), true, $this->maxBurst));
  67. } else {
  68. $remainingTokens = $tokens - $availableTokens;
  69. $waitDuration = $this->rate->calculateTimeForTokens($remainingTokens);
  70. if (null !== $maxTime && $waitDuration > $maxTime) {
  71. // process needs to wait longer than set interval
  72. $rateLimit = new RateLimit($availableTokens, \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false, $this->maxBurst);
  73. throw new MaxWaitDurationExceededException(sprintf('The rate limiter wait time ("%d" seconds) is longer than the provided maximum time ("%d" seconds).', $waitDuration, $maxTime), $rateLimit);
  74. }
  75. // at $now + $waitDuration all tokens will be reserved for this process,
  76. // so no tokens are left for other processes.
  77. $bucket->setTokens($availableTokens - $tokens);
  78. $bucket->setTimer($now);
  79. $reservation = new Reservation($now + $waitDuration, new RateLimit(0, \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false, $this->maxBurst));
  80. }
  81. $this->storage->save($bucket);
  82. } finally {
  83. $this->lock->release();
  84. }
  85. return $reservation;
  86. }
  87. /**
  88. * {@inheritdoc}
  89. */
  90. public function consume(int $tokens = 1): RateLimit
  91. {
  92. try {
  93. return $this->reserve($tokens, 0)->getRateLimit();
  94. } catch (MaxWaitDurationExceededException $e) {
  95. return $e->getRateLimit();
  96. }
  97. }
  98. }