Merge pull request #5 from AgentSoftware/NOTICKET-rate-limit

Improved rate limit logic
AgentSoftware · Apr 8, 2024 · c4d81c6 · c4d81c6
2 parents bd56d31 + e958138
commit c4d81c6
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 41 deletions.
diff --git a/composer.json b/composer.json
@@ -11,8 +11,9 @@
     "require": {
         "php": ">=8.1",
         "ext-json": "*",
-        "illuminate/queue": "^8 || ^9 || ^10",
-        "aws/aws-sdk-php": "^3.2"
+        "aws/aws-sdk-php": "^3.2",
+        "illuminate/cache": "^8 || ^9 || ^10",
+        "illuminate/queue": "^8 || ^9 || ^10"
     },
     "autoload": {
         "psr-4": {

diff --git a/phpstan-baseline.neon b/phpstan-baseline.neon
@@ -20,6 +20,11 @@ parameters:
 			count: 1
 			path: src/RawSqsJob.php
 
+		-
+			message: "#^Cannot call method attempt\\(\\) on mixed\\.$#"
+			count: 1
+			path: src/RawSqsQueue.php
+
 		-
 			message: "#^Method AgentSoftware\\\\LaravelRawSqsConnector\\\\RawSqsQueue\\:\\:log\\(\\) has parameter \\$context with no value type specified in iterable type array\\.$#"
 			count: 1
@@ -41,16 +46,16 @@ parameters:
 			path: src/RawSqsQueue.php
 
 		-
-			message: "#^Method AgentSoftware\\\\LaravelRawSqsConnector\\\\RawSqsQueue\\:\\:querySqs\\(\\) return type has no value type specified in iterable type array\\|Aws\\\\Result\\.$#"
+			message: "#^Method AgentSoftware\\\\LaravelRawSqsConnector\\\\RawSqsQueue\\:\\:receiveMessage\\(\\) return type has no value type specified in iterable type Aws\\\\Result\\.$#"
 			count: 1
 			path: src/RawSqsQueue.php
 
 		-
-			message: "#^Method AgentSoftware\\\\LaravelRawSqsConnector\\\\RawSqsQueue\\:\\:receiveMessage\\(\\) return type has no value type specified in iterable type Aws\\\\Result\\.$#"
+			message: "#^Method AgentSoftware\\\\LaravelRawSqsConnector\\\\RawSqsQueue\\:\\:receiveMessage\\(\\) return type has no value type specified in iterable type array\\.$#"
 			count: 1
 			path: src/RawSqsQueue.php
 
 		-
-			message: "#^Method AgentSoftware\\\\LaravelRawSqsConnector\\\\RawSqsQueue\\:\\:receiveMessage\\(\\) return type has no value type specified in iterable type array\\.$#"
+			message: "#^Method AgentSoftware\\\\LaravelRawSqsConnector\\\\RawSqsQueue\\:\\:receiveMessage\\(\\) should return array\\|Aws\\\\Result\\|null but returns mixed\\.$#"
 			count: 1
 			path: src/RawSqsQueue.php
diff --git a/src/RawSqsQueue.php b/src/RawSqsQueue.php
@@ -8,8 +8,8 @@
 use Illuminate\Queue\Jobs\SqsJob;
 use Illuminate\Queue\SqsQueue;
 use Illuminate\Support\Facades\Log;
-use Illuminate\Support\Facades\RateLimiter;
 use Illuminate\Support\Str;
+use Illuminate\Cache\RateLimiter;
 
 class RawSqsQueue extends SqsQueue
 {
@@ -54,10 +54,10 @@ protected function receiveMessage(string $queue): Result|array|null
 
         $key = 'sqs:' . Str::slug($this->jobClass);
 
-        $remainingAttempts = $this->hasRemainingAttempts($key);
+        $attempt = $this->attempt($key, $queue);
 
-        if ($remainingAttempts) {
-            return $this->querySqs($queue);
+        if ($attempt !== false) {
+            return $attempt;
         }
 
         $this->log('Rate limit hit for SQS queue worker', [
@@ -73,19 +73,24 @@ protected function log(string $text, array $context = []): void
         Log::info($text, $context);
     }
 
-    protected function hasRemainingAttempts(string $key): mixed
+    protected function attempt(string $key, string $queue): mixed
     {
         /** @var int $limit */
         $limit = $this->getRateLimit();
 
-        return RateLimiter::attempt(
+        return $this->getRateLimiter()->attempt(
             $key,
             $limit,
-            fn () => true,
+            fn () => $this->querySqs($queue),
         );
     }
 
-    protected function querySqs(string $queue): Result|array
+    protected function getRateLimiter(): mixed
+    {
+        return $this->getContainer()->make(RateLimiter::class);
+    }
+
+    protected function querySqs(string $queue): array|Result|null
     {
         return $this->sqs->receiveMessage([
             'QueueUrl' => $queue,

diff --git a/tests/RawSqsQueueTest.php b/tests/RawSqsQueueTest.php
@@ -5,11 +5,11 @@
 use Aws\Sqs\SqsClient;
 use Illuminate\Container\Container;
 use Illuminate\Queue\InvalidPayloadException;
-use Illuminate\Support\Facades\RateLimiter;
 use Mockery;
 use PHPUnit\Framework\TestCase;
 use AgentSoftware\LaravelRawSqsConnector\RawSqsQueue;
 use Tests\Support\TestJobClass;
+use Illuminate\Cache\RateLimiter;
 
 class RawSqsQueueTest extends TestCase
 {
@@ -136,14 +136,15 @@ public function testDoesNotUseRateLimiterIfRateLimitNotSpecified(): void
         ];
 
         $sqsClientMock = Mockery::mock(SqsClient::class);
+
         $sqsClientMock->shouldReceive('receiveMessage')
             ->andReturn([
                 'Messages' => [
                     $sqsReturnMessage
                 ]
             ]);
 
-        $sqsClientMock->shouldNotReceive('hasRemainingAttempts');
+        $sqsClientMock->shouldNotReceive('attempt');
 
         $rawSqsQueue = new RawSqsQueue(
             $sqsClientMock,
@@ -173,26 +174,29 @@ public function testWillReturnMessageIfRateLimitEnabled(): void
         ];
 
         $sqsClientMock = Mockery::mock(SqsClient::class);
-        $sqsClientMock->shouldReceive('receiveMessage')
+
+        $rawSqsQueue = new RawSqsQueue(
+            $sqsClientMock,
+            'default',
+            'prefix'
+        );
+
+        $rateLimiter = Mockery::mock(RateLimiter::class);
+
+        $rateLimiter->shouldReceive('attempt')
+            ->once()
             ->andReturn([
                 'Messages' => [
                     $sqsReturnMessage
                 ]
             ]);
 
-        $rawSqsQueue = Mockery::mock(RawSqsQueue::class, [
-            $sqsClientMock,
-            'default',
-            'prefix'
-        ])
-            ->shouldAllowMockingProtectedMethods()
-            ->makePartial();
+        $container = Mockery::mock(Container::class);
 
-        $rawSqsQueue
-            ->shouldReceive('hasRemainingAttempts')
-            ->andReturn(true);
+        $container->shouldReceive('make')
+            ->once()
+            ->andReturn($rateLimiter);
 
-        $container = Mockery::mock(Container::class);
         $rawSqsQueue->setContainer($container);
         $rawSqsQueue->setJobClass(TestJobClass::class);
         $rawSqsQueue->setRateLimit(1);
@@ -202,33 +206,32 @@ public function testWillReturnMessageIfRateLimitEnabled(): void
         $this->expectNotToPerformAssertions();
     }
 
-    public function testWillNotReturnMessageIfRateLimitEnabledButNoAttemptsLeft(): void
+    public function testWillNotReturnMessageIfRateLimitHit(): void
     {
         $sqsClientMock = Mockery::mock(SqsClient::class);
-        $sqsClientMock->shouldNotReceive('receiveMessage');
 
         $rawSqsQueue = Mockery::mock(RawSqsQueue::class, [
             $sqsClientMock,
             'default',
             'prefix'
-        ])
-            ->makePartial();
+        ])->makePartial();
 
-        $rawSqsQueue
-            ->shouldAllowMockingProtectedMethods()
-            ->shouldReceive('hasRemainingAttempts')
-            ->andReturn(false);
-
-        $rawSqsQueue
-            ->shouldAllowMockingProtectedMethods()
+        $rawSqsQueue->shouldAllowMockingProtectedMethods()
             ->shouldReceive('log')
             ->once();
 
-        $rawSqsQueue
-            ->shouldAllowMockingProtectedMethods()
-            ->shouldNotReceive('querySqs');
+        $rateLimiter = Mockery::mock(RateLimiter::class);
+
+        $rateLimiter->shouldReceive('attempt')
+            ->once()
+            ->andReturn(false);
 
         $container = Mockery::mock(Container::class);
+
+        $container->shouldReceive('make')
+            ->once()
+            ->andReturn($rateLimiter);
+
         $rawSqsQueue->setContainer($container);
         $rawSqsQueue->setJobClass(TestJobClass::class);
         $rawSqsQueue->setRateLimit(1);