8000 [RateLimiter] Return Limit object on Consume method · symfony/symfony@8f62afc · GitHub
[go: up one dir, main page]
Skip to content

Commit 8f62afc

Browse files
Valentinfabpot
Valentin
authored and
fabpot
committed
[RateLimiter] Return Limit object on Consume method
1 parent a429dee commit 8f62afc

File tree

12 files changed

+128
-31
lines changed

12 files changed

+128
-31
lines changed

‎src/Symfony/Component/RateLimiter/CompoundLimiter.php

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,17 +25,28 @@ final class CompoundLimiter implements LimiterInterface
2525
*/
2626
public function __construct(array $limiters)
2727
{
28+
if (!$limiters) {
29+
throw new \LogicException(sprintf('"%s::%s()" require at least one limiter.', self::class, __METHOD__));
30+
}
2831
$this->limiters = $limiters;
2932
}
3033

31-
public function consume(int $tokens = 1): bool
34+
public function consume(int $tokens = 1): Limit
3235
{
33-
$allow = true;
36+
$minimalLimit = null;
3437
foreach ($this->limiters as $limiter) {
35-
$allow = $limiter->consume($tokens) && $allow;
38+
$limit = $limiter->consume($tokens);
39+
40+
if (0 === $limit->getRemainingTokens()) {
41+
return $limit;
42+
}
43+
44+
if (null === $minimalLimit || $limit->getRemainingTokens() < $minimalLimit->getRemainingTokens()) {
45+
$minimalLimit = $limit;
46+
}
3647
}
3748

38-
return $allow;
49+
return $minimalLimit;
3950
}
4051

4152
public function reset(): void

‎src/Symfony/Component/RateLimiter/FixedWindowLimiter.php

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ public function __construct(string $id, int $limit, \DateInterval $interval, Sto
4343
/**
4444
* {@inheritdoc}
4545
*/
46-
public function consume(int $tokens = 1): bool
46+
public function consume(int $tokens = 1): Limit
4747
{
4848
$this->lock->acquire(true);
4949

@@ -54,17 +54,28 @@ public function consume(int $tokens = 1): bool
5454
}
5555

5656
$hitCount = $window->getHitCount();
57-
$availableTokens = $this->limit - $hitCount;
57+
$availableTokens = $this->getAvailableTokens($hitCount);
58+
$windowStart = \DateTimeImmutable::createFromFormat('U', time());
5859
if ($availableTokens < $tokens) {
59-
return false;
60+
return new Limit($availableTokens, $this->getRetryAfter($windowStart), false);
6061
}
6162

6263
$window->add($tokens);
6364
$this->storage->save($window);
6465

65-
return true;
66+
return new Limit($this->getAvailableTokens($window->getHitCount()), $this->getRetryAfter($windowStart), true);
6667
} finally {
6768
$this->lock->release();
6869
}
6970
}
71+
72+
public function getAvailableTokens(int $hitCount): int
73+
{
74+
return $this->limit - $hitCount;
75+
}
76+
77+
private function getRetryAfter(\DateTimeImmutable $windowStart): \DateTimeImmutable
78+
{
79+
return $windowStart->add(new \DateInterval(sprintf('PT%sS', $this->interval)));
80+
}
7081
}

‎src/Symfony/Component/RateLimiter/Limit.php

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
<?php
2+
3+
/*
4+
* This file is part of the Symfony package.
5+
*
6+
* (c) Fabien Potencier <fabien@symfony.com>
7+
*
8+
* For the full copyright and license information, please view the LICENSE
9+
* file that was distributed with this source code.
10+
*/
11+
12+
namespace Symfony\Component\RateLimiter;
13+
14+
/**
15+
* @author Valentin Silvestre <vsilvestre.pro@gmail.com>
16+
*
17+
* @experimental in 5.2
18+
*/
19+
class Limit
20+
{
21+
private $availableTokens;
22+
private $retryAfter;
23+
private $accepted;
24+
25+
public function __construct(int $availableTokens, \DateTimeImmutable $retryAfter, bool $accepted)
26+
{
27+
$this->availableTokens = $availableTokens;
28+
$this->retryAfter = $retryAfter;
29+
$this->accepted = $accepted;
30+
}
31+
32+
public function isAccepted(): bool
33+
{
34+
return $this->accepted;
35+
}
36+
37+
public function getRetryAfter(): \DateTimeImmutable
38+
{
39+
return $this->retryAfter;
40+
}
41+
42+
public function getRemainingTokens(): int
43+
{
44+
return $this->availableTokens;
45+
}
46+
}

‎src/Symfony/Component/RateLimiter/LimiterInterface.php

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ interface LimiterInterface
2424
*
2525
* @param int $tokens the number of tokens required
2626
*/
27-
public function consume(int $tokens = 1): bool;
27+
public function consume(int $tokens = 1): Limit;
2828

2929
/**
3030
* Resets the limit.

‎src/Symfony/Component/RateLimiter/NoLimiter.php

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,9 @@
2323
*/
2424
final class NoLimiter implements LimiterInterface
2525
{
26-
public function consume(int $tokens = 1): bool
26+
public function consume(int $tokens = 1): Limit
2727
{
28-
return true;
28+
return new Limit(\INF, new \DateTimeImmutable(), true, 'no_limit');
2929
}
3030

3131
public function reset(): void

‎src/Symfony/Component/RateLimiter/README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ $limiter->reserve(1)->wait();
3232
// ... execute the code
3333

3434
// only claims 1 token if it's free at this moment (useful if you plan to skip this process)
35-
if ($limiter->consume(1)) {
35+
if ($limiter->consume(1)->isAccepted()) {
3636
// ... execute the code
3737
}
3838
```

‎src/Symfony/Component/RateLimiter/Rate.php

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,16 @@ public function calculateTimeForTokens(int $tokens): int
7373
return TimeUtil::dateIntervalToSeconds($this->refillTime) * $cyclesRequired;
7474
}
7575

76+
/**
77+
* Calculates the next moment of token availability.
78+
*
79+
* @return \DateTimeImmutable the next moment a token will be available
80+
*/
81+
public function calculateNextTokenAvailability(): \DateTimeImmutable
82+
{
83+
return (new \DateTimeImmutable())->add($this->refillTime);
84+
}
85+
7686
/**
7787
* Calculates the number of new free tokens during $duration.
7888
*

‎src/Symfony/Component/RateLimiter/Tests/CompoundLimiterTest.php

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -38,19 +38,20 @@ public function testConsume()
3838
$limiter3 = $this->createLimiter(12, new \DateInterval('PT30S'));
3939
$limiter = new CompoundLimiter([$limiter1, $limiter2, $limiter3]);
4040

41-
$this->assertFalse($limiter->consume(5), 'Limiter 1 reached the limit');
41+
// Reach limiter 1 limit, verify that limiter2 available tokens reduced by 5 and fetch successfully limiter 1
42+
$this->assertEquals(3, $limiter->consume(5)->getRemainingTokens(), 'Limiter 1 reached the limit');
4243
sleep(1); // reset limiter1's window
43-
$limiter->consume(2);
44+
$this->assertTrue($limiter->consume(2)->isAccepted());
4445

45-
$this->assertTrue($limiter->consume());
46-
$this->assertFalse($limiter->consume(), 'Limiter 2 reached the limit');
46+
// Reach limiter 2 limit, verify that limiter2 available tokens reduced by 5 and and fetch successfully
47+
$this->assertEquals(0, $limiter->consume()->getRemainingTokens(), 'Limiter 2 has no remaining tokens left');
4748
sleep(9); // reset limiter2's window
49+
$this->assertTrue($limiter->consume(3)->isAccepted());
4850

49-
$this->assertTrue($limiter->consume(3));
50-
$this->assertFalse($limiter->consume(), 'Limiter 3 reached the limit');
51+
// Reach limiter 3 limit, verify that limiter2 available tokens reduced by 5 and fetch successfully
52+
$this->assertEquals(0, $limiter->consume()->getRemainingTokens(), 'Limiter 3 reached the limit');
5153
sleep(20); // reset limiter3's window
52-
53-
$this->assertTrue($limiter->consume());
54+
$this->assertTrue($limiter->consume()->isAccepted());
5455
}
5556

5657
private function createLimiter(int $limit, \DateInterval $interval): FixedWindowLimiter

‎src/Symfony/Component/RateLimiter/Tests/FixedWindowLimiterTest.php

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,8 +40,10 @@ public function testConsume()
4040
sleep(5);
4141
}
4242

43-
$this->assertTrue($limiter->consume());
44-
$this->assertFalse($limiter->consume());
43+
$limit = $limiter->consume();
44+
$this->assertTrue($limit->isAccepted());
45+
$limit = $limiter->consume();
46+
$this->assertFalse($limit->isAccepted());
4547
}
4648

4749
public function testConsumeOutsideInterval()
@@ -55,7 +57,9 @@ public function testConsumeOutsideInterval()
5557
$limiter->consume(9);
5658
// ...try bursting again at the start of the next window
5759
sleep(10);
58-
$this->assertTrue($limiter->consume(10));
60+
$limit = $limiter->consume(10);
61+
$this->assertEquals(0, $limit->getRemainingTokens());
62+
$this->assertEquals(time() + 60, $limit->getRetryAfter()->getTimestamp());
5963
}
6064

6165
private function createLimiter(): FixedWindowLimiter

‎src/Symfony/Component/RateLimiter/Tests/TokenBucketLimiterTest.php

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -69,13 +69,21 @@ public function testReserveMaxWaitingTime()
6969

7070
public function testConsume()
7171
{
72-
$limiter = $this->createLimiter();
72+
$rate = Rate::perSecond(10);
73+
$limiter = $this->createLimiter(10, $rate);
7374

7475
// enough free tokens
75-
$this->assertTrue($limiter->consume(5));
76+
$limit = $limiter->consume(5);
77+
$this->assertTrue($limit->isAccepted());
78+
$this->assertEquals(5, $limit->getRemainingTokens());
79+
$this->assertEqualsWithDelta(time(), $limit->getRetryAfter()->getTimestamp(), 1);
7680
// there are only 5 available free tokens left now
77-
$this->assertFalse($limiter->consume(10));
78-
$this->assertTrue($limiter->consume(5));
81+
$limit = $limiter->consume(10);
82+
$this->assertEquals(5, $limit->getRemainingTokens());
83+
84+
$limit = $limiter->consume(5);
85+
$this->assertEquals(0, $limit->getRemainingTokens());
86+
$this->assertEqualsWithDelta(time(), $limit->getRetryAfter()->getTimestamp(), 1);
7987
}
8088

8189
private function createLimiter($initialTokens = 10, Rate $rate = null)

‎src/Symfony/Component/RateLimiter/TokenBucketLimiter.php

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -103,14 +103,20 @@ public function reserve(int $tokens = 1, ?float $maxTime = null): Reservation
103103
/**
104104
* {@inheritdoc}
105105
*/
106-
public function consume(int $tokens = 1): bool
106+
public function consume(int $tokens = 1): Limit
107107
{
108+
$bucket = $this->storage->fetch($this->id);
109+
if (null === $bucket) {
110+
$bucket = new TokenBucket($this->id, $this->maxBurst, $this->rate);
111+
}
112+
$now = microtime(true);
113+
108114
try {
109115
$this->reserve($tokens, 0);
110116

111-
return true;
117+
return new Limit($bucket->getAvailableTokens($now) - $tokens, $this->rate->calculateNextTokenAvailability(), true);
112118
} catch (MaxWaitDurationExceededException $e) {
113-
return false;
119+
return new Limit($bucket->getAvailableTokens($now), $this->rate->calculateNextTokenAvailability(), false);
114120
}
115121
}
116122
}

‎src/Symfony/Component/Security/Http/EventListener/LoginThrottlingListener.php

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ public function checkPassport(CheckPassportEvent $event): void
4848
$limiterKey = $this->createLimiterKey($username, $request);
4949

5050
$limiter = $this->limiter->create($limiterKey);
51-
if (!$limiter->consume()) {
51+
if (!$limiter->consume()->isAccepted()) {
5252
throw new TooManyLoginAttemptsAuthenticationException();
5353
}
5454
}

0 commit comments

Comments
 (0)
0