diff --git a/src/Aeon/RateLimiter/Algorithm.php b/src/Aeon/RateLimiter/Algorithm.php index 19d0d7b..adc2164 100644 --- a/src/Aeon/RateLimiter/Algorithm.php +++ b/src/Aeon/RateLimiter/Algorithm.php @@ -17,5 +17,10 @@ public function hit(string $id, Storage $storage) : void; /** * Estimate the time in which next hit is allowed. */ - public function nextHit(string $id, Storage $storage) : TimeUnit; + public function estimate(string $id, Storage $storage) : TimeUnit; + + /** + * Return hits left before throttling next hit. + */ + public function capacity(string $id, Storage $storage) : int; } diff --git a/src/Aeon/RateLimiter/Algorithm/LeakyBucketAlgorithm.php b/src/Aeon/RateLimiter/Algorithm/LeakyBucketAlgorithm.php index 2d4b539..4c48f80 100644 --- a/src/Aeon/RateLimiter/Algorithm/LeakyBucketAlgorithm.php +++ b/src/Aeon/RateLimiter/Algorithm/LeakyBucketAlgorithm.php @@ -48,7 +48,7 @@ public function hit(string $id, Storage $storage) : void /** * @psalm-suppress PossiblyNullReference */ - public function nextHit(string $id, Storage $storage) : TimeUnit + public function estimate(string $id, Storage $storage) : TimeUnit { $hits = $storage->all($id); @@ -59,4 +59,11 @@ public function nextHit(string $id, Storage $storage) : TimeUnit return TimeUnit::seconds(0); } + + public function capacity(string $id, Storage $storage) : int + { + $hits = $storage->all($id); + + return $this->bucketSize - $hits->count(); + } } diff --git a/src/Aeon/RateLimiter/Algorithm/SlidingWindowAlgorithm.php b/src/Aeon/RateLimiter/Algorithm/SlidingWindowAlgorithm.php index a8bff32..bda69bf 100644 --- a/src/Aeon/RateLimiter/Algorithm/SlidingWindowAlgorithm.php +++ b/src/Aeon/RateLimiter/Algorithm/SlidingWindowAlgorithm.php @@ -43,7 +43,7 @@ public function hit(string $id, Storage $storage) : void /** * @psalm-suppress PossiblyNullReference */ - public function nextHit(string $id, Storage $storage) : TimeUnit + public function estimate(string $id, Storage $storage) : TimeUnit { $hits = $storage->all($id); @@ -54,4 +54,11 @@ public function nextHit(string $id, Storage $storage) : TimeUnit return TimeUnit::seconds(0); } + + public function capacity(string $id, Storage $storage) : int + { + $hits = $storage->all($id); + + return $this->limit - $hits->count(); + } } diff --git a/src/Aeon/RateLimiter/RateLimiter.php b/src/Aeon/RateLimiter/RateLimiter.php index 0021556..b2f7f56 100644 --- a/src/Aeon/RateLimiter/RateLimiter.php +++ b/src/Aeon/RateLimiter/RateLimiter.php @@ -30,7 +30,12 @@ public function hit(string $id) : void public function estimate(string $id) : TimeUnit { - return $this->algorithm->nextHit($id, $this->storage); + return $this->algorithm->estimate($id, $this->storage); + } + + public function capacity(string $id) : int + { + return $this->algorithm->capacity($id, $this->storage); } public function throttle(string $id, Process $process) : void diff --git a/tests/Aeon/RateLimiter/Tests/Unit/Algorithm/LeakyBucketAlgorithmTest.php b/tests/Aeon/RateLimiter/Tests/Unit/Algorithm/LeakyBucketAlgorithmTest.php index adaf76c..7a76382 100644 --- a/tests/Aeon/RateLimiter/Tests/Unit/Algorithm/LeakyBucketAlgorithmTest.php +++ b/tests/Aeon/RateLimiter/Tests/Unit/Algorithm/LeakyBucketAlgorithmTest.php @@ -22,22 +22,26 @@ public function test_leaky_bucket_algorithm() : void $algorithm = new LeakyBucketAlgorithm($calendar, $bucketSize = 5, $leakSize = 2, TimeUnit::seconds(10)); - $algorithm->hit('id', $storage = new MemoryStorage($calendar)); + $this->assertSame($bucketSize, $algorithm->capacity('id', $storage = new MemoryStorage($calendar))); + + $algorithm->hit('id', $storage); $algorithm->hit('id', $storage); $algorithm->hit('id', $storage); $algorithm->hit('id', $storage); $algorithm->hit('id', $storage); - $this->assertSame(10, $algorithm->nextHit('id', $storage)->inSeconds()); + $this->assertSame(10, $algorithm->estimate('id', $storage)->inSeconds()); + $this->assertSame(0, $algorithm->capacity('id', $storage)); $calendar->setNow($calendar->now()->add(TimeUnit::seconds(10)->add(TimeUnit::millisecond()))); - $this->assertSame(0, $algorithm->nextHit('id', $storage)->inSeconds()); + $this->assertSame(0, $algorithm->estimate('id', $storage)->inSeconds()); + $this->assertSame(2, $algorithm->capacity('id', $storage)); $algorithm->hit('id', $storage); $algorithm->hit('id', $storage); - $this->assertSame('9.999000', $algorithm->nextHit('id', $storage)->inSecondsPrecise()); + $this->assertSame('9.999000', $algorithm->estimate('id', $storage)->inSecondsPrecise()); } public function test_leaky_bucket_algorithm_with_too_many_hits() : void diff --git a/tests/Aeon/RateLimiter/Tests/Unit/Algorithm/SlidingWindowAlgorithmTest.php b/tests/Aeon/RateLimiter/Tests/Unit/Algorithm/SlidingWindowAlgorithmTest.php index 2a19fba..eba8f74 100644 --- a/tests/Aeon/RateLimiter/Tests/Unit/Algorithm/SlidingWindowAlgorithmTest.php +++ b/tests/Aeon/RateLimiter/Tests/Unit/Algorithm/SlidingWindowAlgorithmTest.php @@ -20,7 +20,7 @@ public function test_with_available_hits() : void $algorithm = new SlidingWindowAlgorithm($calendar = new GregorianCalendarStub(TimeZone::UTC()), 2, TimeUnit::minute()); $algorithm->hit('hit_id', $memoryStorage = new MemoryStorage($calendar)); - $this->assertSame($algorithm->nextHit('hit_id', $memoryStorage)->inSeconds(), 0); + $this->assertSame($algorithm->estimate('hit_id', $memoryStorage)->inSeconds(), 0); } public function test_without_available_hits() : void @@ -28,7 +28,7 @@ public function test_without_available_hits() : void $algorithm = new SlidingWindowAlgorithm($calendar = new GregorianCalendarStub(TimeZone::UTC()), 1, TimeUnit::minute()); $algorithm->hit('hit_id', $memoryStorage = new MemoryStorage($calendar)); - $this->assertSame($algorithm->nextHit('hit_id', $memoryStorage)->inSeconds(), 59); + $this->assertSame($algorithm->estimate('hit_id', $memoryStorage)->inSeconds(), 59); } public function test_hit_without_available_hits() : void @@ -66,22 +66,27 @@ public function test_resetting_hits() : void $calendar->setNow(DateTime::fromString('2020-01-01 00:00:00 UTC')); - $algorithm->hit('hit_id', $memoryStorage = new MemoryStorage($calendar)); + $this->assertSame(1, $algorithm->capacity('hit_id', $storage = new MemoryStorage($calendar))); + $algorithm->hit('hit_id', $storage); + $this->assertSame(0, $algorithm->capacity('hit_id', $storage)); - $this->assertSame($algorithm->nextHit('hit_id', $memoryStorage)->inSeconds(), 60); + $this->assertSame($algorithm->estimate('hit_id', $storage)->inSeconds(), 60); $calendar->setNow($calendar->now()->add(TimeUnit::seconds(61))); - $this->assertSame($algorithm->nextHit('hit_id', $memoryStorage)->inSeconds(), 0); + $this->assertSame(1, $algorithm->capacity('hit_id', $storage)); + $this->assertSame($algorithm->estimate('hit_id', $storage)->inSeconds(), 0); - $algorithm->hit('hit_id', $memoryStorage); + $algorithm->hit('hit_id', $storage); - $this->assertSame($algorithm->nextHit('hit_id', $memoryStorage)->inSeconds(), 60); + $this->assertSame(0, $algorithm->capacity('hit_id', $storage)); + $this->assertSame($algorithm->estimate('hit_id', $storage)->inSeconds(), 60); $calendar->setNow($calendar->now()->add(TimeUnit::seconds(61))); - $this->assertSame($algorithm->nextHit('hit_id', $memoryStorage)->inSeconds(), 0); + $this->assertSame(1, $algorithm->capacity('hit_id', $storage)); + $this->assertSame($algorithm->estimate('hit_id', $storage)->inSeconds(), 0); - $algorithm->hit('hit_id', $memoryStorage); + $algorithm->hit('hit_id', $storage); } } diff --git a/tests/Aeon/RateLimiter/Tests/Unit/RateLimiterTest.php b/tests/Aeon/RateLimiter/Tests/Unit/RateLimiterTest.php index 4c20469..98a2ecd 100644 --- a/tests/Aeon/RateLimiter/Tests/Unit/RateLimiterTest.php +++ b/tests/Aeon/RateLimiter/Tests/Unit/RateLimiterTest.php @@ -33,7 +33,7 @@ public function test_hit_method_throwing_exception() : void public function test_estimate_method_throwing_exception() : void { $algorithm = $this->createStub(Algorithm::class); - $algorithm->method('nextHit')->willReturn(TimeUnit::second()); + $algorithm->method('estimate')->willReturn(TimeUnit::second()); $rateLimiter = new RateLimiter( $algorithm, @@ -58,6 +58,19 @@ public function test_throttle_without_waiting() : void $rateLimiter->throttle('id', $process); } + public function test_capacity() : void + { + $algorithm = $this->createStub(Algorithm::class); + $algorithm->method('capacity')->willReturn(10); + + $rateLimiter = new RateLimiter( + $algorithm, + $this->createMock(Storage::class) + ); + + $this->assertSame(10, $rateLimiter->capacity('id')); + } + public function test_throttle_and_wait() : void { $algorithm = $this->createMock(Algorithm::class);