use lua script to pop the task atomically

Author: michalsn

src/Handlers/PredisHandler.php

@@ -13,6 +13,7 @@
 
 namespace CodeIgniter\Queue\Handlers;
 
+use CodeIgniter\Autoloader\FileLocator;
 use CodeIgniter\Exceptions\CriticalError;
 use CodeIgniter\I18n\Time;
 use CodeIgniter\Queue\Config\Queue as QueueConfig;
@@ -27,12 +28,20 @@
 class PredisHandler extends BaseHandler implements QueueInterface
 {
     private readonly Client $predis;
+    private readonly string $luaScript;
 
     public function __construct(protected QueueConfig $config)
     {
         try {
             $this->predis = new Client($config->predis, ['prefix' => $config->predis['prefix']]);
             $this->predis->time();
+
+            $locator   = new FileLocator(service('autoloader'));
+            $luaScript = $locator->locateFile('CodeIgniter\Queue\Lua\pop_task', null, 'lua');
+            if ($luaScript === false) {
+                throw new CriticalError('Queue: LUA script for Predis is not available.');
+            }
+            $this->luaScript = file_get_contents($luaScript);
         } catch (Exception $e) {
             throw new CriticalError('Queue: Predis connection refused (' . $e->getMessage() . ').');
         }
@@ -77,30 +86,23 @@ public function push(string $queue, string $job, array $data): bool
      */
     public function pop(string $queue, array $priorities): ?QueueJob
     {
-        $tasks = [];
-        $now   = Time::now()->timestamp;
-
-        foreach ($priorities as $priority) {
-            $tasks = $this->predis->zrangebyscore(
-                "queues:{$queue}:{$priority}",
-                '-inf',
-                $now,
-                ['LIMIT' => [0, 1]]
-            );
-            if ($tasks !== []) {
-                $removed = $this->predis->zrem("queues:{$queue}:{$priority}", ...$tasks);
-                if ($removed !== 0) {
-                    break;
-                }
-                $tasks = [];
-            }
-        }
+        $now = Time::now()->timestamp;
+
+        // Prepare the arguments for the Lua script
+        $args = [
+            'queues:' . $queue,       // KEYS[1]
+            $now,                     // ARGV[2]
+            json_encode($priorities), // ARGV[3]
+        ];
+
+        // Execute the Lua script
+        $task = $this->predis->eval($this->luaScript, 1, ...$args);
 
-        if ($tasks === []) {
+        if ($task === null) {
             return null;
         }
 
-        $queueJob = new QueueJob(json_decode((string) $tasks[0], true));
+        $queueJob = new QueueJob(json_decode((string) $task, true));
 
         // Set the actual status as in DB.
         $queueJob->status = Status::RESERVED->value;

src/Handlers/RedisHandler.php

@@ -13,6 +13,7 @@
 
 namespace CodeIgniter\Queue\Handlers;
 
+use CodeIgniter\Autoloader\FileLocator;
 use CodeIgniter\Exceptions\CriticalError;
 use CodeIgniter\I18n\Time;
 use CodeIgniter\Queue\Config\Queue as QueueConfig;
@@ -27,6 +28,7 @@
 class RedisHandler extends BaseHandler implements QueueInterface
 {
     private readonly Redis $redis;
+    private readonly string $luaScript;
 
     public function __construct(protected QueueConfig $config)
     {
@@ -48,6 +50,13 @@ public function __construct(protected QueueConfig $config)
             if (isset($config->redis['prefix']) && ! $this->redis->setOption(Redis::OPT_PREFIX, $config->redis['prefix'])) {
                 throw new CriticalError('Queue: Redis setting prefix failed.');
             }
+
+            $locator   = new FileLocator(service('autoloader'));
+            $luaScript = $locator->locateFile('CodeIgniter\Queue\Lua\pop_task', null, 'lua');
+            if ($luaScript === false) {
+                throw new CriticalError('Queue: LUA script for Redis is not available.');
+            }
+            $this->luaScript = file_get_contents($luaScript);
         } catch (RedisException $e) {
             throw new CriticalError('Queue: RedisException occurred with message (' . $e->getMessage() . ').');
         }
@@ -96,23 +105,23 @@ public function push(string $queue, string $job, array $data): bool
      */
     public function pop(string $queue, array $priorities): ?QueueJob
     {
-        $tasks = [];
-        $now   = Time::now()->timestamp;
-
-        foreach ($priorities as $priority) {
-            if ($tasks = $this->redis->zRangeByScore("queues:{$queue}:{$priority}", '-inf', (string) $now, ['limit' => [0, 1]])) {
-                if ($this->redis->zRem("queues:{$queue}:{$priority}", ...$tasks)) {
-                    break;
-                }
-                $tasks = [];
-            }
-        }
+        $now = Time::now()->timestamp;
+
+        // Prepare the arguments for the Lua script
+        $args = [
+            'queues:' . $queue,       // KEYS[1]
+            $now,                     // ARGV[2]
+            json_encode($priorities), // ARGV[3]
+        ];
+
+        // Execute the Lua script
+        $task = $this->redis->eval($this->luaScript, $args, 1);
 
-        if ($tasks === []) {
+        if ($task === false) {
             return null;
         }
 
-        $queueJob = new QueueJob(json_decode((string) $tasks[0], true));
+        $queueJob = new QueueJob(json_decode((string) $task, true));
 
         // Set the actual status as in DB.
         $queueJob->status = Status::RESERVED->value;

src/Lua/pop_task.lua

@@ -0,0 +1,17 @@
+local queue = KEYS[1]
+local now = tonumber(ARGV[1])
+local priorities = cjson.decode(ARGV[2])
+local task = nil
+
+for _, priority in ipairs(priorities) do
+    local key = queue .. ':' .. priority
+    local tasks = redis.call('ZRANGEBYSCORE', key, '-inf', tostring(now), 'LIMIT', 0, 1)
+
+    if #tasks > 0 then
+        redis.call('ZREM', key, tasks[1])
+        task = tasks[1]
+        break
+    end
+end
+
+return task