Added method to get current capacity of rate limiter for given algori…

…thm (#1)
aeon-php · Nov 22, 2020 · b1602e7 · b1602e7
1 parent 43494cf
commit b1602e7
Show file tree

Hide file tree

Showing 7 changed files with 64 additions and 18 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -20,15 +20,15 @@ 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
     {
         $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
@@ -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);