255 lines
7.7 KiB
Python
255 lines
7.7 KiB
Python
|
"""Use pika with the Asyncio EventLoop"""
|
||
|
|
import asyncio
|
||
|
|
from functools import partial
|
||
|
|
|
||
|
|
from pika.adapters import base_connection
|
||
|
|
|
||
|
|
|
||
|
|
class IOLoopAdapter:
|
||
|
|
def __init__(self, loop):
|
||
|
|
"""
|
||
|
|
Basic adapter for asyncio event loop
|
||
|
|
|
||
|
|
:type loop: asyncio.AbstractEventLoop
|
||
|
|
:param loop: Asyncio Loop
|
||
|
|
|
||
|
|
"""
|
||
|
|
self.loop = loop
|
||
|
|
|
||
|
|
self.handlers = {}
|
||
|
|
self.readers = set()
|
||
|
|
self.writers = set()
|
||
|
|
|
||
|
|
def close(self):
|
||
|
|
"""Release ioloop's resources.
|
||
|
|
|
||
|
|
This method is intended to be called by the application or test code
|
||
|
|
only after the ioloop's outermost `start()` call returns. After calling
|
||
|
|
`close()`, no other interaction with the closed instance of ioloop
|
||
|
|
should be performed.
|
||
|
|
|
||
|
|
"""
|
||
|
|
self.loop.close()
|
||
|
|
|
||
|
|
def add_timeout(self, deadline, callback_method):
|
||
|
|
"""Add the callback_method to the EventLoop timer to fire after deadline
|
||
|
|
seconds. Returns a Handle to the timeout.
|
||
|
|
|
||
|
|
:param int deadline: The number of seconds to wait to call callback
|
||
|
|
:param method callback_method: The callback method
|
||
|
|
:rtype: asyncio.Handle
|
||
|
|
|
||
|
|
"""
|
||
|
|
return self.loop.call_later(deadline, callback_method)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def remove_timeout(handle):
|
||
|
|
"""
|
||
|
|
Cancel asyncio.Handle
|
||
|
|
|
||
|
|
:type handle: asyncio.Handle
|
||
|
|
:rtype: bool
|
||
|
|
"""
|
||
|
|
return handle.cancel()
|
||
|
|
|
||
|
|
def add_callback_threadsafe(self, callback):
|
||
|
|
"""Requests a call to the given function as soon as possible in the
|
||
|
|
context of this IOLoop's thread.
|
||
|
|
|
||
|
|
NOTE: This is the only thread-safe method offered by the IOLoop adapter.
|
||
|
|
All other manipulations of the IOLoop adapter and its parent connection
|
||
|
|
must be performed from the connection's thread.
|
||
|
|
|
||
|
|
For example, a thread may request a call to the
|
||
|
|
`channel.basic_ack` method of a connection that is running in a
|
||
|
|
different thread via
|
||
|
|
|
||
|
|
```
|
||
|
|
connection.add_callback_threadsafe(
|
||
|
|
functools.partial(channel.basic_ack, delivery_tag=...))
|
||
|
|
```
|
||
|
|
|
||
|
|
:param method callback: The callback method; must be callable.
|
||
|
|
|
||
|
|
"""
|
||
|
|
self.loop.call_soon_threadsafe(callback)
|
||
|
|
|
||
|
|
def add_handler(self, fd, cb, event_state):
|
||
|
|
""" Registers the given handler to receive the given events for ``fd``.
|
||
|
|
|
||
|
|
The ``fd`` argument is an integer file descriptor.
|
||
|
|
|
||
|
|
The ``event_state`` argument is a bitwise or of the constants
|
||
|
|
``base_connection.BaseConnection.READ``, ``base_connection.BaseConnection.WRITE``,
|
||
|
|
and ``base_connection.BaseConnection.ERROR``.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
if fd in self.handlers:
|
||
|
|
raise ValueError("fd {} added twice".format(fd))
|
||
|
|
self.handlers[fd] = cb
|
||
|
|
|
||
|
|
if event_state & base_connection.BaseConnection.READ:
|
||
|
|
self.loop.add_reader(
|
||
|
|
fd,
|
||
|
|
partial(
|
||
|
|
cb,
|
||
|
|
fd=fd,
|
||
|
|
events=base_connection.BaseConnection.READ
|
||
|
|
)
|
||
|
|
)
|
||
|
|
self.readers.add(fd)
|
||
|
|
|
||
|
|
if event_state & base_connection.BaseConnection.WRITE:
|
||
|
|
self.loop.add_writer(
|
||
|
|
fd,
|
||
|
|
partial(
|
||
|
|
cb,
|
||
|
|
fd=fd,
|
||
|
|
events=base_connection.BaseConnection.WRITE
|
||
|
|
)
|
||
|
|
)
|
||
|
|
self.writers.add(fd)
|
||
|
|
|
||
|
|
def remove_handler(self, fd):
|
||
|
|
""" Stop listening for events on ``fd``. """
|
||
|
|
|
||
|
|
if fd not in self.handlers:
|
||
|
|
return
|
||
|
|
|
||
|
|
if fd in self.readers:
|
||
|
|
self.loop.remove_reader(fd)
|
||
|
|
self.readers.remove(fd)
|
||
|
|
|
||
|
|
if fd in self.writers:
|
||
|
|
self.loop.remove_writer(fd)
|
||
|
|
self.writers.remove(fd)
|
||
|
|
|
||
|
|
del self.handlers[fd]
|
||
|
|
|
||
|
|
def update_handler(self, fd, event_state):
|
||
|
|
if event_state & base_connection.BaseConnection.READ:
|
||
|
|
if fd not in self.readers:
|
||
|
|
self.loop.add_reader(
|
||
|
|
fd,
|
||
|
|
partial(
|
||
|
|
self.handlers[fd],
|
||
|
|
fd=fd,
|
||
|
|
events=base_connection.BaseConnection.READ
|
||
|
|
)
|
||
|
|
)
|
||
|
|
self.readers.add(fd)
|
||
|
|
else:
|
||
|
|
if fd in self.readers:
|
||
|
|
self.loop.remove_reader(fd)
|
||
|
|
self.readers.remove(fd)
|
||
|
|
|
||
|
|
if event_state & base_connection.BaseConnection.WRITE:
|
||
|
|
if fd not in self.writers:
|
||
|
|
self.loop.add_writer(
|
||
|
|
fd,
|
||
|
|
partial(
|
||
|
|
self.handlers[fd],
|
||
|
|
fd=fd,
|
||
|
|
events=base_connection.BaseConnection.WRITE
|
||
|
|
)
|
||
|
|
)
|
||
|
|
self.writers.add(fd)
|
||
|
|
else:
|
||
|
|
if fd in self.writers:
|
||
|
|
self.loop.remove_writer(fd)
|
||
|
|
self.writers.remove(fd)
|
||
|
|
|
||
|
|
|
||
|
|
def start(self):
|
||
|
|
""" Start Event Loop """
|
||
|
|
if self.loop.is_running():
|
||
|
|
return
|
||
|
|
|
||
|
|
self.loop.run_forever()
|
||
|
|
|
||
|
|
def stop(self):
|
||
|
|
""" Stop Event Loop """
|
||
|
|
if self.loop.is_closed():
|
||
|
|
return
|
||
|
|
|
||
|
|
self.loop.stop()
|
||
|
|
|
||
|
|
|
||
|
|
class AsyncioConnection(base_connection.BaseConnection):
|
||
|
|
""" The AsyncioConnection runs on the Asyncio EventLoop.
|
||
|
|
|
||
|
|
:param pika.connection.Parameters parameters: Connection parameters
|
||
|
|
:param on_open_callback: The method to call when the connection is open
|
||
|
|
:type on_open_callback: method
|
||
|
|
:param on_open_error_callback: Method to call if the connection cant be opened
|
||
|
|
:type on_open_error_callback: method
|
||
|
|
:param asyncio.AbstractEventLoop loop: By default asyncio.get_event_loop()
|
||
|
|
|
||
|
|
"""
|
||
|
|
def __init__(self,
|
||
|
|
parameters=None,
|
||
|
|
on_open_callback=None,
|
||
|
|
on_open_error_callback=None,
|
||
|
|
on_close_callback=None,
|
||
|
|
stop_ioloop_on_close=False,
|
||
|
|
custom_ioloop=None):
|
||
|
|
""" Create a new instance of the AsyncioConnection class, connecting
|
||
|
|
to RabbitMQ automatically
|
||
|
|
|
||
|
|
:param pika.connection.Parameters parameters: Connection parameters
|
||
|
|
:param on_open_callback: The method to call when the connection is open
|
||
|
|
:type on_open_callback: method
|
||
|
|
:param on_open_error_callback: Method to call if the connection cant be opened
|
||
|
|
:type on_open_error_callback: method
|
||
|
|
:param asyncio.AbstractEventLoop loop: By default asyncio.get_event_loop()
|
||
|
|
|
||
|
|
"""
|
||
|
|
self.sleep_counter = 0
|
||
|
|
self.loop = custom_ioloop or asyncio.get_event_loop()
|
||
|
|
self.ioloop = IOLoopAdapter(self.loop)
|
||
|
|
|
||
|
|
super().__init__(
|
||
|
|
parameters, on_open_callback,
|
||
|
|
on_open_error_callback,
|
||
|
|
on_close_callback, self.ioloop,
|
||
|
|
stop_ioloop_on_close=stop_ioloop_on_close,
|
||
|
|
)
|
||
|
|
|
||
|
|
def _adapter_connect(self):
|
||
|
|
"""Connect to the remote socket, adding the socket to the EventLoop if
|
||
|
|
connected.
|
||
|
|
|
||
|
|
:rtype: bool
|
||
|
|
|
||
|
|
"""
|
||
|
|
error = super()._adapter_connect()
|
||
|
|
|
||
|
|
if not error:
|
||
|
|
self.ioloop.add_handler(
|
||
|
|
self.socket.fileno(),
|
||
|
|
self._handle_events,
|
||
|
|
self.event_state,
|
||
|
|
)
|
||
|
|
|
||
|
|
return error
|
||
|
|
|
||
|
|
def _adapter_disconnect(self):
|
||
|
|
"""Disconnect from the RabbitMQ broker"""
|
||
|
|
|
||
|
|
if self.socket:
|
||
|
|
self.ioloop.remove_handler(
|
||
|
|
self.socket.fileno()
|
||
|
|
)
|
||
|
|
|
||
|
|
super()._adapter_disconnect()
|
||
|
|
|
||
|
|
def _handle_disconnect(self):
|
||
|
|
# No other way to handle exceptions.ProbableAuthenticationError
|
||
|
|
try:
|
||
|
|
super()._handle_disconnect()
|
||
|
|
super()._handle_write()
|
||
|
|
except Exception as e:
|
||
|
|
# FIXME: Pass None or other constant instead "-1"
|
||
|
|
self._on_disconnect(-1, e)
|