Fewer allocations in consumer

composer.json

@@ -35,7 +35,8 @@
             "Thesis\\Pgmq\\": "src/"
         },
         "files": [
-            "src/pgmq.php"
+            "src/pgmq.php",
+            "src/Internal/functions.php"
         ]
     },
     "autoload-dev": {

src/ConsumeController.php

@@ -5,16 +5,20 @@
 namespace Thesis\Pgmq;
 
 use Amp\Postgres\PostgresTransaction;
+use Thesis\Pgmq;
 use Thesis\Time\TimeSpan;
 
 /**
  * @api
  */
 final readonly class ConsumeController
 {
+    /**
+     * @param non-empty-string $queue
+     */
     public function __construct(
         public PostgresTransaction $tx,
-        private Queue $queue,
+        private string $queue,
         private ConsumeContext $context,
     ) {}
 
@@ -28,23 +32,23 @@ public function stop(): void
      */
     public function ack(array $messages): void
     {
-        $this->queue->deleteBatch(array_map(
-            static fn(Message $message): int => $message->id,
-            $messages,
-        ));
+        Pgmq\deleteBatch(
+            pg: $this->tx,
+            queue: $this->queue,
+            messageIds: array_map(static fn(Message $m) => $m->id, $messages),
+        );
     }
 
     /**
      * @param non-empty-list<Message> $messages
      */
     public function nack(array $messages, TimeSpan $delay): void
     {
-        $this->queue->setVisibilityTimeout(
-            array_map(
-                static fn(Message $message): int => $message->id,
-                $messages,
-            ),
-            $delay,
+        Pgmq\setVisibilityTimeout(
+            pg: $this->tx,
+            queue: $this->queue,
+            messageIds: array_map(static fn(Message $m) => $m->id, $messages),
+            visibilityTimeout: $delay,
         );
     }
 
@@ -53,9 +57,10 @@ public function nack(array $messages, TimeSpan $delay): void
      */
     public function term(array $messages): void
     {
-        $this->queue->archiveBatch(array_map(
-            static fn(Message $message): int => $message->id,
-            $messages,
-        ));
+        Pgmq\archiveBatch(
+            pg: $this->tx,
+            queue: $this->queue,
+            messageIds: array_map(static fn(Message $m) => $m->id, $messages),
+        );
     }
 }

src/Consumer.php

@@ -8,13 +8,14 @@
 use Amp\Pipeline;
 use Amp\Postgres\PostgresConnection;
 use function Amp\async;
+use function Thesis\Pgmq\Internal\consume;
 
 /**
  * @api
  */
 final class Consumer
 {
-    /** @var list<Internal\ConsumeHandler> */
+    /** @var list<ConsumeContext> */
     private array $consumers = [];
 
     public function __construct(
@@ -56,17 +57,15 @@ public function consume(
             );
         }
 
-        $handle = new Internal\ConsumeHandler(
+        $context = consume(
             $this->pg,
             $config,
             $handler,
             new Internal\AggregateWatcher($watchers),
             $polls,
         );
 
-        $this->consumers[] = $handle;
-
-        return $handle->context;
+        return $this->consumers[] = $context;
     }
 
     public function stop(): void

src/Internal/functions.php

@@ -0,0 +1,84 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Thesis\Pgmq\Internal;
+
+use Amp\DeferredFuture;
+use Amp\Pipeline;
+use Amp\Postgres\PostgresConnection;
+use Revolt\EventLoop;
+use Thesis\Pgmq;
+
+/**
+ * @param callable(non-empty-list<Pgmq\Message>, Pgmq\ConsumeController): void $handler
+ * @param Pipeline\Queue<null> $polls
+ */
+function consume(
+    PostgresConnection $pg,
+    Pgmq\ConsumeConfig $config,
+    callable $handler,
+    PollWatcher $watcher,
+    Pipeline\Queue $polls,
+): Pgmq\ConsumeContext {
+    $iterator = $polls->iterate();
+    $completionMarker = new DeferredFuture();
+
+    $context = new Pgmq\ConsumeContext(
+        stop: static function () use ($polls): void {
+            if (!$polls->isComplete()) {
+                $polls->complete();
+            }
+        },
+        completionMarker: $completionMarker->getFuture(),
+    );
+
+    EventLoop::queue(static function () use (
+        $pg,
+        $config,
+        $handler,
+        $iterator,
+        $watcher,
+        $completionMarker,
+        $context,
+    ): void {
+        $watcher->watch();
+
+        foreach ($iterator as $_) {
+            $tx = $pg->beginTransaction();
+
+            try {
+                $messages = Pgmq\readBatch(
+                    pg: $tx,
+                    queue: $config->queue,
+                    count: $config->batch,
+                    visibilityTimeout: $config->visibilityTimeout,
+                );
+
+                $messages = [...$messages];
+
+                if (\count($messages) > 0) {
+                    $handler(
+                        $messages,
+                        new Pgmq\ConsumeController($tx, $config->queue, $context),
+                    );
+                }
+
+                $tx->commit();
+            } catch (\Throwable $e) {
+                $tx->rollback();
+                $completionMarker->error($e);
+                $iterator->dispose();
+                break;
+            }
+        }
+
+        $watcher->cancel();
+
+        if (!$completionMarker->isComplete()) {
+            $completionMarker->complete();
+        }
+    });
+
+    return $context;
+}