From 00d4739adc98ec8dd5f659c7f5e40fe8c4ea6c74 Mon Sep 17 00:00:00 2001 From: Paras Malhotra Date: Wed, 28 Oct 2020 20:53:32 +0530 Subject: [PATCH] [8.x] Add dontRelease option to RateLimited and RateLimitedWithRedis job middleware --- .../Queue/Middleware/RateLimited.php | 23 ++++++++++- .../Queue/Middleware/RateLimitedWithRedis.php | 4 +- tests/Integration/Queue/RateLimitedTest.php | 39 ++++++++++++++++++ .../Queue/RateLimitedWithRedisTest.php | 41 +++++++++++++++++++ 4 files changed, 105 insertions(+), 2 deletions(-) diff --git a/src/Illuminate/Queue/Middleware/RateLimited.php b/src/Illuminate/Queue/Middleware/RateLimited.php index dcedc6bff36..62ed3ac61d2 100644 --- a/src/Illuminate/Queue/Middleware/RateLimited.php +++ b/src/Illuminate/Queue/Middleware/RateLimited.php @@ -23,6 +23,13 @@ class RateLimited */ protected $limiterName; + /** + * Indicates if the job should be released if the limit is exceeded. + * + * @var bool + */ + public $shouldRelease = true; + /** * Create a new middleware instance. * @@ -68,6 +75,18 @@ public function handle($job, $next) ); } + /** + * Do not release the job back to the queue if limit is exceeded. + * + * @return $this + */ + public function dontRelease() + { + $this->shouldRelease = false; + + return $this; + } + /** * Handle a rate limited job. * @@ -80,7 +99,9 @@ protected function handleJob($job, $next, array $limits) { foreach ($limits as $limit) { if ($this->limiter->tooManyAttempts($limit->key, $limit->maxAttempts)) { - return $job->release($this->getTimeUntilNextRetry($limit->key)); + return $this->shouldRelease + ? $job->release($this->getTimeUntilNextRetry($limit->key)) + : false; } $this->limiter->hit($limit->key, $limit->decayMinutes * 60); diff --git a/src/Illuminate/Queue/Middleware/RateLimitedWithRedis.php b/src/Illuminate/Queue/Middleware/RateLimitedWithRedis.php index e4916c17380..7000e9d5530 100644 --- a/src/Illuminate/Queue/Middleware/RateLimitedWithRedis.php +++ b/src/Illuminate/Queue/Middleware/RateLimitedWithRedis.php @@ -50,7 +50,9 @@ protected function handleJob($job, $next, array $limits) { foreach ($limits as $limit) { if ($this->tooManyAttempts($limit->key, $limit->maxAttempts, $limit->decayMinutes)) { - return $job->release($this->getTimeUntilNextRetry($limit->key)); + return $this->shouldRelease + ? $job->release($this->getTimeUntilNextRetry($limit->key)) + : false; } } diff --git a/tests/Integration/Queue/RateLimitedTest.php b/tests/Integration/Queue/RateLimitedTest.php index 0877ccf837c..c4f12f6144e 100644 --- a/tests/Integration/Queue/RateLimitedTest.php +++ b/tests/Integration/Queue/RateLimitedTest.php @@ -49,6 +49,18 @@ public function testRateLimitedJobsAreNotExecutedOnLimitReached() $this->assertJobWasReleased(RateLimitedTestJob::class); } + public function testRateLimitedJobsCanBeSkippedOnLimitReached() + { + $rateLimiter = $this->app->make(RateLimiter::class); + + $rateLimiter->for('test', function ($job) { + return Limit::perHour(1); + }); + + $this->assertJobRanSuccessfully(RateLimitedDontReleaseTestJob::class); + $this->assertJobWasSkipped(RateLimitedDontReleaseTestJob::class); + } + public function testJobsCanHaveConditionalRateLimits() { $rateLimiter = $this->app->make(RateLimiter::class); @@ -105,6 +117,25 @@ protected function assertJobWasReleased($class) $this->assertFalse($class::$handled); } + + protected function assertJobWasSkipped($class) + { + $class::$handled = false; + $instance = new CallQueuedHandler(new Dispatcher($this->app), $this->app); + + $job = m::mock(Job::class); + + $job->shouldReceive('hasFailed')->once()->andReturn(false); + $job->shouldReceive('isReleased')->once()->andReturn(false); + $job->shouldReceive('isDeletedOrReleased')->once()->andReturn(false); + $job->shouldReceive('delete')->once(); + + $instance->call($job, [ + 'command' => serialize($command = new $class), + ]); + + $this->assertFalse($class::$handled); + } } class RateLimitedTestJob @@ -139,3 +170,11 @@ public function isAdmin() return false; } } + +class RateLimitedDontReleaseTestJob extends RateLimitedTestJob +{ + public function middleware() + { + return [(new RateLimited('test'))->dontRelease()]; + } +} diff --git a/tests/Integration/Queue/RateLimitedWithRedisTest.php b/tests/Integration/Queue/RateLimitedWithRedisTest.php index 1be53bfc101..5956fff0618 100644 --- a/tests/Integration/Queue/RateLimitedWithRedisTest.php +++ b/tests/Integration/Queue/RateLimitedWithRedisTest.php @@ -66,6 +66,20 @@ public function testRateLimitedJobsAreNotExecutedOnLimitReached() $this->assertJobWasReleased($testJob); } + public function testRateLimitedJobsCanBeSkippedOnLimitReached() + { + $rateLimiter = $this->app->make(RateLimiter::class); + + $testJob = new RedisRateLimitedDontReleaseTestJob; + + $rateLimiter->for($testJob->key, function ($job) { + return Limit::perMinute(1); + }); + + $this->assertJobRanSuccessfully($testJob); + $this->assertJobWasSkipped($testJob); + } + public function testJobsCanHaveConditionalRateLimits() { $rateLimiter = $this->app->make(RateLimiter::class); @@ -134,6 +148,25 @@ protected function assertJobWasReleased($testJob) $this->assertFalse($testJob::$handled); } + + protected function assertJobWasSkipped($testJob) + { + $testJob::$handled = false; + $instance = new CallQueuedHandler(new Dispatcher($this->app), $this->app); + + $job = m::mock(Job::class); + + $job->shouldReceive('hasFailed')->once()->andReturn(false); + $job->shouldReceive('isReleased')->once()->andReturn(false); + $job->shouldReceive('isDeletedOrReleased')->once()->andReturn(false); + $job->shouldReceive('delete')->once(); + + $instance->call($job, [ + 'command' => serialize($testJob), + ]); + + $this->assertFalse($testJob::$handled); + } } class RedisRateLimitedTestJob @@ -175,3 +208,11 @@ public function isAdmin() return false; } } + +class RedisRateLimitedDontReleaseTestJob extends RedisRateLimitedTestJob +{ + public function middleware() + { + return [(new RateLimitedWithRedis($this->key))->dontRelease()]; + } +}