From ac6348abe206bf8dd3daae96d3d0f20242fd1200 Mon Sep 17 00:00:00 2001 From: Evgeny Ruban Date: Wed, 13 Dec 2023 20:30:40 +0400 Subject: [PATCH] [RateLimiter] Fix results on last token consume. --- .../RateLimiter/Policy/FixedWindowLimiter.php | 22 ++++++++++++--- .../RateLimiter/Policy/TokenBucketLimiter.php | 24 ++++++++++++---- .../Tests/Policy/FixedWindowLimiterTest.php | 24 ++++++++++++++++ .../Tests/Policy/TokenBucketLimiterTest.php | 28 +++++++++++++++++++ 4 files changed, 89 insertions(+), 9 deletions(-) diff --git a/src/Symfony/Component/RateLimiter/Policy/FixedWindowLimiter.php b/src/Symfony/Component/RateLimiter/Policy/FixedWindowLimiter.php index 298cd1ba24321..7c9fe4d0c254f 100644 --- a/src/Symfony/Component/RateLimiter/Policy/FixedWindowLimiter.php +++ b/src/Symfony/Component/RateLimiter/Policy/FixedWindowLimiter.php @@ -63,21 +63,35 @@ public function reserve(int $tokens = 1, float $maxTime = null): Reservation $now = microtime(true); $availableTokens = $window->getAvailableTokens($now); - if ($availableTokens >= $tokens) { + if (0 !== $tokens && $availableTokens > $tokens) { $window->add($tokens, $now); $reservation = new Reservation($now, new RateLimit($window->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now)), true, $this->limit)); } else { + if ($availableTokens === $tokens) { + $window->add($tokens, $now); + } + $waitDuration = $window->calculateTimeForTokens($tokens, $now); - if (null !== $maxTime && $waitDuration > $maxTime) { + if ($availableTokens !== $tokens && 0 !== $tokens && 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($window->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false, $this->limit)); } - $window->add($tokens, $now); + if ($availableTokens !== $tokens) { + $window->add($tokens, $now); + } + + if ($availableTokens === $tokens || 0 === $tokens) { + $accepted = true; + $timeToAct = $now; + } else { + $accepted = false; + $timeToAct = $now + $waitDuration; + } - $reservation = new Reservation($now + $waitDuration, new RateLimit($window->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false, $this->limit)); + $reservation = new Reservation($timeToAct, new RateLimit($window->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), $accepted, $this->limit)); } $this->storage->save($window); } finally { diff --git a/src/Symfony/Component/RateLimiter/Policy/TokenBucketLimiter.php b/src/Symfony/Component/RateLimiter/Policy/TokenBucketLimiter.php index dd00346671362..e84001c87de8f 100644 --- a/src/Symfony/Component/RateLimiter/Policy/TokenBucketLimiter.php +++ b/src/Symfony/Component/RateLimiter/Policy/TokenBucketLimiter.php @@ -67,16 +67,22 @@ public function reserve(int $tokens = 1, float $maxTime = null): Reservation $now = microtime(true); $availableTokens = $bucket->getAvailableTokens($now); - if ($availableTokens >= $tokens) { + if (0 !== $tokens && $availableTokens > $tokens) { // tokens are now available, update bucket $bucket->setTokens($availableTokens - $tokens); $reservation = new Reservation($now, new RateLimit($bucket->getAvailableTokens($now), \DateTimeImmutable::createFromFormat('U', floor($now)), true, $this->maxBurst)); } else { - $remainingTokens = $tokens - $availableTokens; - $waitDuration = $this->rate->calculateTimeForTokens($remainingTokens); + if ($availableTokens === $tokens) { + $waitDuration = $this->rate->calculateTimeForTokens($tokens); + } elseif (0 === $tokens) { + $waitDuration = 0; + } else { + $remainingTokens = $tokens - $availableTokens; + $waitDuration = $this->rate->calculateTimeForTokens($remainingTokens); + } - if (null !== $maxTime && $waitDuration > $maxTime) { + if ($availableTokens !== $tokens && 0 !== $tokens && null !== $maxTime && $waitDuration > $maxTime) { // process needs to wait longer than set interval $rateLimit = new RateLimit($availableTokens, \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false, $this->maxBurst); @@ -87,7 +93,15 @@ public function reserve(int $tokens = 1, float $maxTime = null): Reservation // so no tokens are left for other processes. $bucket->setTokens($availableTokens - $tokens); - $reservation = new Reservation($now + $waitDuration, new RateLimit(0, \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), false, $this->maxBurst)); + if ($availableTokens === $tokens || 0 === $tokens) { + $accepted = true; + $timeToAct = $now; + } else { + $accepted = false; + $timeToAct = $now + $waitDuration; + } + + $reservation = new Reservation($timeToAct, new RateLimit(0, \DateTimeImmutable::createFromFormat('U', floor($now + $waitDuration)), $accepted, $this->maxBurst)); } $this->storage->save($bucket); diff --git a/src/Symfony/Component/RateLimiter/Tests/Policy/FixedWindowLimiterTest.php b/src/Symfony/Component/RateLimiter/Tests/Policy/FixedWindowLimiterTest.php index 84df7bc850aae..f3ea64df82ef2 100644 --- a/src/Symfony/Component/RateLimiter/Tests/Policy/FixedWindowLimiterTest.php +++ b/src/Symfony/Component/RateLimiter/Tests/Policy/FixedWindowLimiterTest.php @@ -57,6 +57,21 @@ public function testConsume() $this->assertEquals($retryAfter, $rateLimit->getRetryAfter()); } + public function testConsumeLastToken() + { + $now = time(); + $limiter = $this->createLimiter(); + $limiter->consume(9); + + $rateLimit = $limiter->consume(1); + $this->assertSame(0, $rateLimit->getRemainingTokens()); + $this->assertTrue($rateLimit->isAccepted()); + $this->assertEquals( + \DateTimeImmutable::createFromFormat('U', $now + 60), + $rateLimit->getRetryAfter() + ); + } + /** * @dataProvider provideConsumeOutsideInterval */ @@ -76,6 +91,15 @@ public function testConsumeOutsideInterval(string $dateIntervalString) $this->assertTrue($rateLimit->isAccepted()); } + public function testReserve() + { + $limiter = $this->createLimiter('PT1S'); + + $this->assertEquals(0, $limiter->reserve(5)->getWaitDuration()); + $this->assertEquals(0, $limiter->reserve(5)->getWaitDuration()); + $this->assertEquals(1, $limiter->reserve(5)->getWaitDuration()); + } + public function testWaitIntervalOnConsumeOverLimit() { $limiter = $this->createLimiter(); diff --git a/src/Symfony/Component/RateLimiter/Tests/Policy/TokenBucketLimiterTest.php b/src/Symfony/Component/RateLimiter/Tests/Policy/TokenBucketLimiterTest.php index 3b7b579c0cf77..ed22bdd244377 100644 --- a/src/Symfony/Component/RateLimiter/Tests/Policy/TokenBucketLimiterTest.php +++ b/src/Symfony/Component/RateLimiter/Tests/Policy/TokenBucketLimiterTest.php @@ -91,6 +91,34 @@ public function testConsume() $this->assertSame(10, $rateLimit->getLimit()); } + public function testConsumeLastToken() + { + $rate = Rate::perSecond(1); + $limiter = $this->createLimiter(10, $rate); + + $rateLimit = $limiter->consume(10); + $this->assertSame(0, $rateLimit->getRemainingTokens()); + $this->assertTrue($rateLimit->isAccepted()); + $this->assertEqualsWithDelta(time(), $rateLimit->getRetryAfter()->getTimestamp(), 10); + } + + public function testConsumeZeroTokens() + { + $rate = Rate::perSecond(1); + $limiter = $this->createLimiter(10, $rate); + + $rateLimit = $limiter->consume(0); + $this->assertTrue($rateLimit->isAccepted()); + $this->assertEquals(time(), $rateLimit->getRetryAfter()->getTimestamp()); + + $limiter->reset(); + $limiter->consume(10); + + $rateLimit = $limiter->consume(0); + $this->assertTrue($rateLimit->isAccepted()); + $this->assertEquals(time(), $rateLimit->getRetryAfter()->getTimestamp()); + } + public function testWaitIntervalOnConsumeOverLimit() { $limiter = $this->createLimiter();