matekasse/venv/lib/python3.11/site-packages/socketio/asyncio_aiopika_manager.py
2023-07-28 23:30:45 +02:00

126 lines
5.1 KiB
Python

import asyncio
import pickle
from socketio.asyncio_pubsub_manager import AsyncPubSubManager
try:
import aio_pika
except ImportError:
aio_pika = None
class AsyncAioPikaManager(AsyncPubSubManager): # pragma: no cover
"""Client manager that uses aio_pika for inter-process messaging under
asyncio.
This class implements a client manager backend for event sharing across
multiple processes, using RabbitMQ
To use a aio_pika backend, initialize the :class:`Server` instance as
follows::
url = 'amqp://user:password@hostname:port//'
server = socketio.Server(client_manager=socketio.AsyncAioPikaManager(
url))
:param url: The connection URL for the backend messaging queue. Example
connection URLs are ``'amqp://guest:guest@localhost:5672//'``
for RabbitMQ.
:param channel: The channel name on which the server sends and receives
notifications. Must be the same in all the servers.
With this manager, the channel name is the exchange name
in rabbitmq
:param write_only: If set to ``True``, only initialize to emit events. The
default of ``False`` initializes the class for emitting
and receiving.
"""
name = 'asyncaiopika'
def __init__(self, url='amqp://guest:guest@localhost:5672//',
channel='socketio', write_only=False, logger=None):
if aio_pika is None:
raise RuntimeError('aio_pika package is not installed '
'(Run "pip install aio_pika" in your '
'virtualenv).')
self.url = url
self._lock = asyncio.Lock()
self.publisher_connection = None
self.publisher_channel = None
self.publisher_exchange = None
super().__init__(channel=channel, write_only=write_only, logger=logger)
async def _connection(self):
return await aio_pika.connect_robust(self.url)
async def _channel(self, connection):
return await connection.channel()
async def _exchange(self, channel):
return await channel.declare_exchange(self.channel,
aio_pika.ExchangeType.FANOUT)
async def _queue(self, channel, exchange):
queue = await channel.declare_queue(durable=False,
arguments={'x-expires': 300000})
await queue.bind(exchange)
return queue
async def _publish(self, data):
if self.publisher_connection is None:
async with self._lock:
if self.publisher_connection is None:
self.publisher_connection = await self._connection()
self.publisher_channel = await self._channel(
self.publisher_connection
)
self.publisher_exchange = await self._exchange(
self.publisher_channel
)
retry = True
while True:
try:
await self.publisher_exchange.publish(
aio_pika.Message(
body=pickle.dumps(data),
delivery_mode=aio_pika.DeliveryMode.PERSISTENT
), routing_key='*',
)
break
except aio_pika.AMQPException:
if retry:
self._get_logger().error('Cannot publish to rabbitmq... '
'retrying')
retry = False
else:
self._get_logger().error(
'Cannot publish to rabbitmq... giving up')
break
except aio_pika.exceptions.ChannelInvalidStateError:
# aio_pika raises this exception when the task is cancelled
raise asyncio.CancelledError()
async def _listen(self):
async with (await self._connection()) as connection:
channel = await self._channel(connection)
await channel.set_qos(prefetch_count=1)
exchange = await self._exchange(channel)
queue = await self._queue(channel, exchange)
retry_sleep = 1
while True:
try:
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
yield pickle.loads(message.body)
retry_sleep = 1
except aio_pika.AMQPException:
self._get_logger().error(
'Cannot receive from rabbitmq... '
'retrying in {} secs'.format(retry_sleep))
await asyncio.sleep(retry_sleep)
retry_sleep = min(retry_sleep * 2, 60)
except aio_pika.exceptions.ChannelInvalidStateError:
# aio_pika raises this exception when the task is cancelled
raise asyncio.CancelledError()