Feat: add support for concurrency manager

.github/workflows/tests.yml

@@ -42,5 +42,6 @@ jobs:
           docker compose up -d
           sleep 10
 
+      # When re-running locally, always do 'docker compose down -v' beforehand
       - name: Run Tests
         run: docker compose exec tests vendor/bin/phpunit --configuration phpunit.xml

src/Queue/Client.php

@@ -16,6 +16,10 @@ public function __construct(string $queue, Connection $connection, string $names
 
     public function enqueue(array $payload): bool
     {
+        if (!isset($payload['concurrencyKey'])) {
+            $payload['concurrencyKey'] = \uniqid();
+        }
+
         $payload = [
             'pid' => \uniqid(more_entropy: true),
             'queue' => $this->queue,

src/Queue/Concurrency/Manager.php

@@ -0,0 +1,38 @@
+<?php
+
+namespace Utopia\Queue\Concurrency;
+
+use Utopia\Queue\Connection;
+use Utopia\Queue\Message;
+
+abstract class Manager
+{
+    public function __construct(protected Connection $adapter, protected int $limit = 1)
+    {
+    }
+
+    public function canProcessJob(Message $message): bool
+    {
+        $key = $this->getConcurrencyKey($message);
+        $value = $this->adapter->get($key);
+        if ($value === null) {
+            $this->adapter->set($key, "0");
+            $value = 0;
+        }
+        return \intval($value) < $this->limit;
+    }
+
+    public function startJob(Message $message): void
+    {
+        $key = $this->getConcurrencyKey($message);
+        $this->adapter->increment($key);
+    }
+
+    public function finishJob(Message $message): void
+    {
+        $key = $this->getConcurrencyKey($message);
+        $this->adapter->decrement($key);
+    }
+
+    abstract public function getConcurrencyKey(Message $message): string;
+}

src/Queue/Worker.php

@@ -7,6 +7,7 @@
 use Utopia\CLI\Console;
 use Utopia\DI\Container;
 use Utopia\DI\Dependency;
+use Utopia\Queue\Concurrency\Manager;
 use Utopia\Servers\Base;
 
 class Worker extends Base
@@ -25,6 +26,13 @@ class Worker extends Base
      */
     protected static Job $job;
 
+    /**
+     * Concurrency Manager
+     *
+     * @var ?Manager
+     */
+    protected ?Manager $concurrencyManager = null;
+
     /**
      * Creates an instance of a Queue server.
      * @param Adapter $adapter
@@ -43,6 +51,27 @@ public static function job(): Job
         return self::$job;
     }
 
+    /**
+     * Register a concurrency manager
+     *
+     * @param Manager $manager
+     */
+    public function setConcurrencyManager(Manager $manager): self
+    {
+        $this->concurrencyManager = $manager;
+        return $this;
+    }
+
+    /**
+     * Get a concurrency manager
+     *
+     * @return ?Manager
+     */
+    public function getConcurrencyManager(): ?Manager
+    {
+        return $this->concurrencyManager;
+    }
+
     /**
      * Stops the Queue server.
      * @return self
@@ -150,10 +179,23 @@ protected function lifecycle(Job $job, Message $message, array $nextMessage, Con
     {
         Console::info("[Job] Received Job ({$message->getPid()}).");
 
+        $concurrencyManager = $this->getConcurrencyManager();
+
         $groups = $job->getGroups();
 
         $connection->getConnection();
 
+        if ($concurrencyManager && !$concurrencyManager->canProcessJob($message)) {
+            $this->adapter->connection->leftPushArray("{$this->adapter->namespace}.queue.{$this->adapter->queue}", $nextMessage);
+            Console::info("[Job] Re-queued Job ({$message->getPid()}) due to concurrency limit.");
+            return $this;
+        }
+
+        /** Increment the concurrency counter */
+        if ($concurrencyManager) {
+            $concurrencyManager->startJob($message);
+        }
+
         /**
          * Move Job to Jobs and it's PID to the processing list.
          */
@@ -257,6 +299,11 @@ protected function lifecycle(Job $job, Message $message, array $nextMessage, Con
                 }
             }
         } finally {
+            /** Decrement the counter so another job can be picked up again */
+            if ($concurrencyManager) {
+                $concurrencyManager->finishJob($message);
+            }
+
             /**
              * Remove Job from Processing.
              */

tests/Queue/E2E/Adapter/Base.php

@@ -152,4 +152,28 @@ public function testRetry(): void
         $this->assertEquals(2, $client->countFailedJobs());
         $this->assertEquals(0, $client->countSuccessfulJobs());
     }
+
+    public function testConcurrencyManager(): void
+    {
+        run(function () {
+            $client = $this->getClient();
+            go(function () use ($client) {
+                $client->resetStats();
+
+                for ($i = 0; $i < 10; $i++) {
+                    $this->assertTrue($client->enqueue([
+                        'type' => 'test_sleep',
+                        'id' => $i
+                    ]));
+                }
+
+                sleep(1);
+
+                $this->assertEquals(10, $client->countTotalJobs());
+                $this->assertEquals(10, $client->countProcessingJobs());
+                $this->assertEquals(0, $client->countFailedJobs());
+                $this->assertEquals(0, $client->countSuccessfulJobs());
+            });
+        });
+    }
 }

tests/Queue/servers/Swoole/worker.php

@@ -5,6 +5,7 @@
 
 use Utopia\DI\Container;
 use Utopia\Queue;
+use Utopia\Queue\Concurrency\Manager;
 use Utopia\Queue\Message;
 use Utopia\Queue\Worker;
 use Utopia\Servers\Validator;
@@ -61,10 +62,13 @@ public function isValid(mixed $value): bool
     }
 }
 
+
+
 $container = new Container();
 $connection = new Queue\Adapter\Swoole\Redis('redis');
 $adapter = new Queue\Adapter\Swoole\Server($connection, 1, 'swoole');
 $server = new Queue\Worker($adapter);
+$server->setConcurrencyManager(new TestConcurrencyManager($connection, 5));
 $server->setContainer($container);
 
 // Server::init()

tests/Queue/servers/tests.php

@@ -1,6 +1,8 @@
 <?php
 
 use Utopia\Queue;
+use Utopia\Queue\Concurrency\Manager;
+use Utopia\Queue\Message;
 
 function handleRequest(Queue\Message $job): void
 {
@@ -48,6 +50,17 @@ function handleRequest(Queue\Message $job): void
         case 'test_exception':
             assert(false);
 
+            // no break
+        case 'test_sleep':
+            sleep(5);
             break;
     }
 }
+
+class TestConcurrencyManager extends Manager
+{
+    public function getConcurrencyKey(Message $message): string
+    {
+        return $message->getPayload()['concurrencyKey'];
+    }
+}