Create and use a Threadsafe pika connection

It wraps a pika SelectConnection with just the APIs that I need, implemented as thread-safe, converting from callbacks to something more ergonomic when dealing with threads.

Author: ryanhiebert

qio/broker_test.py

@@ -74,20 +74,20 @@ def receive_messages():
             receiver = broker.receive(queuespec)
             for message in receiver:
                 received_messages.append(message)
-                # Don't call start() - this should block after prefetch_limit messages
 
         thread = threading.Thread(target=receive_messages)
         thread.start()
 
         # Wait for receiver to reach prefetch limit or timeout
-        thread.join(timeout=1.0)
+        thread.join(timeout=0.1)
 
         # Should have received exactly prefetch_limit messages and be blocked
         assert len(received_messages) == prefetch_limit
         assert thread.is_alive()  # Thread should still be alive (blocked)
 
         broker.shutdown()
         thread.join(timeout=1.0)  # Clean up thread
+        assert not thread.is_alive()
 
     @pytest.mark.timeout(2)
     def test_suspend_resume_affects_prefetch_capacity(self, broker):

qio/pika/broker.py

@@ -1,13 +1,13 @@
 from threading import Lock
 
-from pika import BlockingConnection
-from pika import ConnectionParameters
 from pika import URLParameters
+from pika.connection import Parameters
 
 from qio.broker import Broker
 from qio.queuespec import QueueSpec
 
 from .receiver import PikaReceiver
+from .threadsafe import ThreadsafeConnection
 
 
 class PikaBroker(Broker):
@@ -19,38 +19,23 @@ def from_uri(cls, uri: str, /):
         amqp_uri = "amqp:" + uri.removeprefix("pika:")
         return cls(URLParameters(amqp_uri))
 
-    def __init__(self, connection_params: ConnectionParameters | URLParameters):
-        self.__connection_params = connection_params
-        self.__producer_channel_lock = Lock()
-        self.__producer_channel = BlockingConnection(self.__connection_params).channel()
+    def __init__(self, connection_params: Parameters):
+        self.__connection = ThreadsafeConnection(connection_params)
+        self.__channel = self.__connection.channel()
         self.__shutdown_lock = Lock()
         self.__shutdown = False
         self.__receivers = set[PikaReceiver]()
 
     def enqueue(self, body: bytes, /, *, queue: str):
-        with self.__producer_channel_lock:
-            self.__producer_channel.basic_publish(
-                exchange="",
-                routing_key=queue,
-                body=body,
-            )
+        self.__channel.declare_queue(queue=queue, durable=True)
+        self.__channel.publish(exchange="", routing_key=queue, body=body)
 
     def purge(self, *, queue: str):
-        with self.__producer_channel_lock:
-            self.__producer_channel.queue_declare(queue=queue, durable=True)
-            self.__producer_channel.queue_purge(queue=queue)
+        self.__channel.declare_queue(queue=queue, durable=True)
+        self.__channel.purge(queue=queue)
 
     def receive(self, queuespec: QueueSpec, /) -> PikaReceiver:
-        if not queuespec.queues:
-            raise ValueError("Must specify at least one queue")
-        if len(queuespec.queues) != 1:
-            raise ValueError("Only one queue is supported")
-
-        receiver = PikaReceiver(
-            connection_params=self.__connection_params,
-            queue=queuespec.queues[0],
-            prefetch=queuespec.concurrency,
-        )
+        receiver = PikaReceiver(self.__connection, queuespec)
         self.__receivers.add(receiver)
         return receiver
 
@@ -62,3 +47,7 @@ def shutdown(self):
             self.__shutdown = True
             for receiver in self.__receivers:
                 receiver.shutdown()
+            self.__connection.close()
+
+    def __del__(self):
+        self.shutdown()

qio/pika/receiver.py

@@ -2,34 +2,42 @@
 from threading import Lock
 from typing import cast
 
-from pika import BlockingConnection
-from pika import ConnectionParameters
-from pika import URLParameters
-
 from qio.message import Message
+from qio.queuespec import QueueSpec
 from qio.receiver import Receiver
 
+from .threadsafe import ThreadsafeConnection
+
 
 class PikaReceiver(Receiver):
     def __init__(
         self,
-        *,
-        connection_params: ConnectionParameters | URLParameters,
-        queue: str,
-        prefetch: int,
+        connection: ThreadsafeConnection,
+        queuespec: QueueSpec,
+        /,
     ):
-        self.__connection = BlockingConnection(connection_params)
-        self.__channel = self.__connection.channel()
-        self.__channel.queue_declare(queue=queue, durable=True)
-        self.__prefetch_lock = Lock()
-        self.__prefetch = prefetch
-        self.__channel.basic_qos(prefetch_count=prefetch, global_qos=True)
-        self.__iterator = self.__channel.consume(queue=queue)
+        if len(queuespec.queues) == 0:
+            raise ValueError("Must specify at least one queue")
+        if len(queuespec.queues) != 1:
+            raise ValueError("Only one queue is supported")
+
+        self.__channel = connection.channel()
+        self.__consumer_tag = dict[str, str]()
         self.__tag = dict[Message, int]()
-        self.__suspended = set[Message]()
+
+        self.__channel.declare_queue(queue=queuespec.queues[0], durable=True)
+        self.__prefetch_lock = Lock()
+        self.__prefetch = 0
+        self.__adjust_prefetch(+queuespec.concurrency)
+        self.__channel.consume(queuespec.queues[0])
+
+    def __adjust_prefetch(self, change: int) -> None:
+        with self.__prefetch_lock:
+            self.__prefetch += change
+            self.__channel.qos(prefetch_count=self.__prefetch, global_qos=True)
 
     def __iter__(self) -> Iterator[Message]:
-        for method, _, body in self.__iterator:
+        for method, _, body in self.__channel.messages():
             message = Message(body)
             tag = cast(int, method.delivery_tag)
             self.__tag[message] = tag
@@ -41,39 +49,24 @@ def pause(self, message: Message, /):
         The message processing is not completed, and is expected to unpause,
         but its assigned capacity may be allocated elsewhere temporarily.
         """
-        with self.__prefetch_lock:
-            self.__prefetch += 1
-            prefetch = self.__prefetch  # Memo for the lambda
-            self.__connection.add_callback_threadsafe(
-                lambda: self.__channel.basic_qos(prefetch_count=prefetch)
-            )
+        self.__adjust_prefetch(+1)
 
     def unpause(self, message: Message, /):
         """Unpause processing of a message.
 
         The previously paused message processing is resuming, so its assigned
         capacity is no longer available for allocation elsewhere.
         """
-        with self.__prefetch_lock:
-            self.__prefetch -= 1
-            prefetch = self.__prefetch  # Memo for the lambda
-            self.__connection.add_callback_threadsafe(
-                lambda: self.__channel.basic_qos(prefetch_count=prefetch)
-            )
+        self.__adjust_prefetch(-1)
 
     def finish(self, message: Message, /):
         """Finish processing a message.
 
         The message is done processing, and its assigned capacity may be
         allocated elsewhere permanently.
         """
-        self.__connection.add_callback_threadsafe(
-            lambda: self.__channel.basic_ack(delivery_tag=self.__tag.pop(message))
-        )
+        self.__channel.ack(delivery_tag=self.__tag.pop(message))
 
     def shutdown(self):
-        self.__connection.add_callback_threadsafe(self.__shutdown)
-
-    def __shutdown(self):
-        self.__channel.cancel()
-        self.__connection.close()
+        for consumer_tag in list(self.__consumer_tag.values()):
+            self.__channel.cancel(consumer_tag=consumer_tag)

qio/pika/threadsafe.py

@@ -0,0 +1,198 @@
+from collections.abc import Callable
+from collections.abc import Iterator
+from collections.abc import Mapping
+from concurrent.futures import Future
+from contextlib import suppress
+from queue import Queue
+from queue import ShutDown
+from threading import Event
+from threading import Thread
+from typing import Any
+from typing import AnyStr
+from typing import cast
+
+from pika import SelectConnection
+from pika import frame
+from pika import spec
+from pika.channel import Channel
+from pika.connection import Parameters
+from pika.exceptions import ConnectionClosedByClient
+
+
+class ThreadsafeConnection:
+    def __init__(self, connection_params: Parameters):
+        opened = Future()
+        self.__closed = Future()
+        self.__connection = SelectConnection(
+            connection_params,
+            lambda _: opened.set_result(None),
+            lambda _, exc: opened.set_exception(cast(BaseException, exc)),
+            lambda _, exc: self.__closed.set_exception(exc),
+        )
+        self.__thread = Thread(target=self.__connection.ioloop.start)
+        self.__thread.start()
+        opened.result()
+
+    def __wait[T](self, fn: Callable[[], T]):
+        # TODO: Throw connection exceptions on all waiters
+        event = Event()
+        self.__connection.ioloop.add_callback_threadsafe(
+            lambda: event.set() if fn() else event.set()
+        )
+        event.wait()
+
+    def channel(self, channel_number: int | None = None) -> ThreadsafeChannel:
+        future = Future[Channel]()
+        self.__wait(
+            lambda: self.__connection.channel(
+                channel_number=channel_number,
+                on_open_callback=future.set_result,
+            )
+        )
+        return ThreadsafeChannel(self.__wait, future.result())
+
+    def close(self, reply_code: int = 200, reply_text: str = "Normal shutdown"):
+        self.__wait(
+            lambda: self.__connection.close(
+                reply_code=reply_code,
+                reply_text=reply_text,
+            )
+        )
+        with suppress(ConnectionClosedByClient):
+            self.__closed.result()
+        self.__connection.ioloop.stop()
+        self.__thread.join()
+
+
+class ThreadsafeChannel:
+    def __init__(
+        self,
+        wait: Callable[[Callable[[], Any]], None],
+        channel: Channel,
+    ):
+        self.__wait = wait
+        self.__channel = channel
+        self.__messages = Queue[
+            tuple[spec.Basic.Deliver, spec.BasicProperties, bytes]
+        ]()
+        self.__channel.add_on_close_callback(
+            lambda c, e: self.__messages.shutdown(immediate=True)
+        )
+
+    def declare_queue(
+        self,
+        queue: str,
+        passive: bool = False,
+        durable: bool = False,
+        exclusive: bool = False,
+        auto_delete: bool = False,
+        arguments: Mapping[str, Any] | None = None,
+    ) -> frame.Method[spec.Queue.DeclareOk]:
+        future: Future[frame.Method[spec.Queue.DeclareOk]] = Future()
+        self.__wait(
+            lambda: self.__channel.queue_declare(
+                queue=queue,
+                passive=passive,
+                durable=durable,
+                exclusive=exclusive,
+                auto_delete=auto_delete,
+                arguments=arguments,
+                callback=future.set_result,
+            )
+        )
+        return future.result()
+
+    def publish(
+        self,
+        exchange: str,
+        routing_key: str,
+        body: AnyStr,
+        properties: spec.BasicProperties | None = None,
+        mandatory: bool = False,
+    ):
+        self.__wait(
+            lambda: self.__channel.basic_publish(
+                exchange=exchange,
+                routing_key=routing_key,
+                body=body,
+                properties=properties,
+                mandatory=mandatory,
+            )
+        )
+
+    def purge(self, queue: str) -> frame.Method[spec.Queue.PurgeOk]:
+        future: Future[frame.Method[spec.Queue.PurgeOk]] = Future()
+        self.__wait(
+            lambda: self.__channel.queue_purge(
+                queue=queue,
+                callback=future.set_result,
+            )
+        )
+
+        return future.result()
+
+    def consume(
+        self,
+        queue: str,
+        auto_ack: bool = False,
+        exclusive: bool = False,
+        consumer_tag: str | None = None,
+        arguments: Mapping[str, Any] | None = None,
+    ) -> frame.Method[spec.Basic.ConsumeOk]:
+        future: Future[frame.Method[spec.Basic.ConsumeOk]] = Future()
+        self.__wait(
+            lambda: self.__channel.basic_consume(
+                queue=queue,
+                on_message_callback=lambda _, m, p, b: self.__messages.put((m, p, b)),
+                auto_ack=auto_ack,
+                exclusive=exclusive,
+                consumer_tag=consumer_tag,
+                arguments=arguments,
+                callback=future.set_result,
+            )
+        )
+        return future.result()
+
+    def cancel(self, consumer_tag: str) -> frame.Method[spec.Basic.CancelOk]:
+        future: Future[frame.Method[spec.Basic.CancelOk]] = Future()
+        self.__wait(
+            lambda: self.__channel.basic_cancel(
+                consumer_tag=consumer_tag,
+                callback=lambda r: future.set_result(r),
+            )
+        )
+        return future.result()
+
+    def qos(
+        self,
+        prefetch_size: int = 0,
+        prefetch_count: int = 0,
+        global_qos: bool = False,
+    ) -> frame.Method[spec.Basic.QosOk]:
+        future: Future[frame.Method[spec.Basic.QosOk]] = Future()
+        self.__wait(
+            lambda: self.__channel.basic_qos(
+                prefetch_size=prefetch_size,
+                prefetch_count=prefetch_count,
+                global_qos=global_qos,
+                callback=future.set_result,
+            )
+        )
+        return future.result()
+
+    def ack(self, delivery_tag: int = 0, multiple: bool = False):
+        self.__wait(
+            lambda: self.__channel.basic_ack(
+                delivery_tag=delivery_tag,
+                multiple=multiple,
+            )
+        )
+
+    def messages(
+        self,
+    ) -> Iterator[tuple[spec.Basic.Deliver, spec.BasicProperties, bytes]]:
+        while True:
+            try:
+                yield self.__messages.get()
+            except ShutDown:
+                return