Add token price and maxTime parameters (#4)

* Add token price and maxTime parameters --------- Co-authored-by: Maxim Solovev <[email protected]>
php-http · Jun 14, 2024 · 725165c · 725165c
1 parent 93ead07
commit 725165c
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 2 deletions.
diff --git a/src/ThrottlePlugin.php b/src/ThrottlePlugin.php
@@ -6,21 +6,39 @@
 
 use Http\Client\Common\Plugin;
 use Http\Promise\Promise;
+use InvalidArgumentException;
 use Psr\Http\Message\RequestInterface;
+use Symfony\Component\RateLimiter\Exception\MaxWaitDurationExceededException;
+use Symfony\Component\RateLimiter\Exception\ReserveNotSupportedException;
 use Symfony\Component\RateLimiter\LimiterInterface;
 
 final class ThrottlePlugin implements Plugin
 {
     private LimiterInterface $rateLimiter;
 
-    public function __construct(LimiterInterface $rateLimiter)
+    private int $tokens;
+
+    private ?float $maxTime;
+
+    /**
+     * @param int $tokens the number of tokens required
+     * @param float|null $maxTime maximum accepted waiting time in seconds
+     */
+    public function __construct(LimiterInterface $rateLimiter, int $tokens = 1, ?float $maxTime = null)
     {
         $this->rateLimiter = $rateLimiter;
+        $this->tokens = $tokens;
+        $this->maxTime = $maxTime;
     }
 
+    /**
+     * @throws MaxWaitDurationExceededException if $maxTime is set and the process needs to wait longer than its value (in seconds)
+     * @throws ReserveNotSupportedException if this limiter implementation doesn't support reserving tokens
+     * @throws InvalidArgumentException if $tokens is larger than the maximum burst size
+     */
     public function handleRequest(RequestInterface $request, callable $next, callable $first): Promise
     {
-        $this->rateLimiter->reserve()->wait();
+        $this->rateLimiter->reserve($this->tokens, $this->maxTime)->wait();
 
         return $next($request);
     }

diff --git a/tests/ThrottlePluginTest.php b/tests/ThrottlePluginTest.php
@@ -11,6 +11,7 @@
 use Nyholm\Psr7\Request;
 use PHPUnit\Framework\TestCase;
 use Symfony\Bridge\PhpUnit\ClockMock;
+use Symfony\Component\RateLimiter\Exception\MaxWaitDurationExceededException;
 use Symfony\Component\RateLimiter\RateLimit;
 use Symfony\Component\RateLimiter\RateLimiterFactory;
 use Symfony\Component\RateLimiter\Storage\InMemoryStorage;
@@ -56,4 +57,51 @@ public function testThrottle(): void
         $this->client->sendRequest(new Request('GET', ''));
         $this->assertEqualsWithDelta($timeAfterThrottle, time(), 1);
     }
+
+    public function testTokens(): void
+    {
+        $this->client = new PluginClient($this->mockClient, [
+            new ThrottlePlugin(
+                (new RateLimiterFactory(
+                    ['id' => 'foo', 'policy' => 'fixed_window', 'limit' => 2, 'interval' => '3 seconds'],
+                    new InMemoryStorage(),
+                ))->create(),
+                2,
+            ),
+        ]);
+
+        $time = time();
+        $this->client->sendRequest(new Request('GET', ''));
+        $this->client->sendRequest(new Request('GET', ''));
+        $this->assertEqualsWithDelta($time, ($timeAfterThrottle = time()) - 3, 1);
+
+        $this->client->sendRequest(new Request('GET', ''));
+        $this->assertEqualsWithDelta($timeAfterThrottle, time(), 1);
+    }
+
+    public function testMaxTime(): void
+    {
+        $this->client = new PluginClient($this->mockClient, [
+            new ThrottlePlugin(
+                $rateLimit = (new RateLimiterFactory(
+                    ['id' => 'foo', 'policy' => 'fixed_window', 'limit' => 2, 'interval' => '3 seconds'],
+                    new InMemoryStorage(),
+                ))->create(),
+                1,
+                1,
+            ),
+        ]);
+
+        $this->expectException(MaxWaitDurationExceededException::class);
+        $this->expectExceptionMessage('The rate limiter wait time ("3" seconds) is longer than the provided maximum time ("1" seconds).');
+
+        $time = time();
+        $this->client->sendRequest(new Request('GET', ''));
+        $this->client->sendRequest(new Request('GET', ''));
+        $this->client->sendRequest(new Request('GET', ''));
+        $this->assertEqualsWithDelta($time, ($timeAfterThrottle = time()) - 3, 1);
+
+        $this->client->sendRequest(new Request('GET', ''));
+        $this->assertEqualsWithDelta($timeAfterThrottle, time(), 1);
+    }
 }