""" Contains functionality for working with RabbitMQ """
import asyncio
import pickle
import aio_pika
from aiormq.exceptions import ChannelInvalidStateError
[docs]class QueueHelper:
""" Helper class for exchanging messages between elements of a weasel-pipeline toolchain """
_queues = dict()
_connection = None
_channel = None
[docs] @staticmethod
async def create(amqp_uri: str, prefetch_count: int = 5000):
"""
Create a QueueHelper instance and establish connection to the RabbitMQ server.
:param amqp_uri: The connection URI for the RabbitMQ server
:param prefetch_count:
Controls the number of messages which are pre-fetched from each queue.
Keeping this number high enough will significantly improve message throughput
:return: The created QueueHelper instance
"""
queue_helper = QueueHelper()
await queue_helper.initialize_pika(amqp_uri, prefetch_count)
return queue_helper
[docs] async def initialize_pika(self, amqp_uri: str, prefetch_count: int) -> None:
"""
Initialize the connection to RabbitMQ.
This is only required if the instance is not created via the `create` method.
:param amqp_uri: The connection URI for the RabbitMQ server
:param prefetch_count:
Controls the number of messages which are pre-fetched from each queue.
Keeping this number high enough will significantly improve message throughput
"""
self._connection = await aio_pika.connect_robust(
amqp_uri, loop=asyncio.get_event_loop()
)
self._channel = await self._connection.channel()
await self._channel.set_qos(prefetch_count=prefetch_count)
[docs] async def close(self) -> None:
"""
Close the connection to RabbitMQ.
"""
await self._channel.close()
await self._connection.close()
[docs] async def declare_queue(self, name: str, queue_max_length: int = 1000) -> None:
"""
Declare a queue with the given name.
:param name: The identifying name for the created queue
:param queue_max_length:
Defines the maximum length of this queue. This prevents the queue from growing
indefinitely. When an element calls `put_value` to a queue which is currently full,
this call will be suspended until there is free space in the queue. This way all
elements in the pipeline get restricted to the speed of the bottleneck elements.
"""
queue = await self._channel.declare_queue(
name,
arguments={
"x-max-length": queue_max_length,
"x-overflow": "reject-publish",
},
)
self._queues[name] = queue
[docs] async def put_value(self, name: str, value: object) -> None:
"""
Output a value to the queue of the given name. If the queue is currently full, this
call will be suspended until there is free space again.
:param name: The identifying name of the queue to which the value is written
:param value: The output value. Can be any data type supported by `pickle`
"""
while True:
try:
await self._channel.default_exchange.publish(
aio_pika.Message(body=pickle.dumps(value)), routing_key=name
)
return
except (
aio_pika.exceptions.DeliveryError,
ChannelInvalidStateError,
ConnectionResetError,
ConnectionRefusedError,
RuntimeError,
):
await asyncio.sleep(0.1)
[docs] async def pull_values(self, name: str):
"""
Pull values from the queue of the given name. This will return an iterator that is
supposed to be used in an `async for` loop.
:param name: The identifying name of the queue from which the values are read
"""
async for message in self._queues[name]:
async with message.process():
yield pickle.loads(message.body)