470 lines
16 KiB
Python
470 lines
16 KiB
Python
"""Use pika with the Gevent IOLoop."""
|
|
|
|
import functools
|
|
import logging
|
|
import os
|
|
import threading
|
|
import weakref
|
|
|
|
try:
|
|
import queue
|
|
except ImportError: # Python <= v2.7
|
|
import Queue as queue
|
|
|
|
import gevent
|
|
import gevent.hub
|
|
import gevent.socket
|
|
|
|
import pika.compat
|
|
from pika.adapters.base_connection import BaseConnection
|
|
from pika.adapters.utils.io_services_utils import check_callback_arg
|
|
from pika.adapters.utils.nbio_interface import (
|
|
AbstractIOReference,
|
|
AbstractIOServices,
|
|
)
|
|
from pika.adapters.utils.selector_ioloop_adapter import (
|
|
AbstractSelectorIOLoop,
|
|
SelectorIOServicesAdapter,
|
|
)
|
|
|
|
LOGGER = logging.getLogger(__name__)
|
|
|
|
|
|
class GeventConnection(BaseConnection):
|
|
"""Implementation of pika's ``BaseConnection``.
|
|
|
|
An async selector-based connection which integrates with Gevent.
|
|
"""
|
|
|
|
def __init__(self,
|
|
parameters=None,
|
|
on_open_callback=None,
|
|
on_open_error_callback=None,
|
|
on_close_callback=None,
|
|
custom_ioloop=None,
|
|
internal_connection_workflow=True):
|
|
"""Create a new GeventConnection instance and connect to RabbitMQ on
|
|
Gevent's event-loop.
|
|
|
|
:param pika.connection.Parameters|None parameters: The connection
|
|
parameters
|
|
:param callable|None on_open_callback: The method to call when the
|
|
connection is open
|
|
:param callable|None on_open_error_callback: Called if the connection
|
|
can't be established or connection establishment is interrupted by
|
|
`Connection.close()`:
|
|
on_open_error_callback(Connection, exception)
|
|
:param callable|None on_close_callback: Called when a previously fully
|
|
open connection is closed:
|
|
`on_close_callback(Connection, exception)`, where `exception` is
|
|
either an instance of `exceptions.ConnectionClosed` if closed by
|
|
user or broker or exception of another type that describes the
|
|
cause of connection failure
|
|
:param gevent._interfaces.ILoop|nbio_interface.AbstractIOServices|None
|
|
custom_ioloop: Use a custom Gevent ILoop.
|
|
:param bool internal_connection_workflow: True for autonomous connection
|
|
establishment which is default; False for externally-managed
|
|
connection workflow via the `create_connection()` factory
|
|
"""
|
|
if pika.compat.ON_WINDOWS:
|
|
raise RuntimeError('GeventConnection is not supported on Windows.')
|
|
|
|
custom_ioloop = (custom_ioloop or
|
|
_GeventSelectorIOLoop(gevent.get_hub()))
|
|
|
|
if isinstance(custom_ioloop, AbstractIOServices):
|
|
nbio = custom_ioloop
|
|
else:
|
|
nbio = _GeventSelectorIOServicesAdapter(custom_ioloop)
|
|
|
|
super(GeventConnection, self).__init__(
|
|
parameters,
|
|
on_open_callback,
|
|
on_open_error_callback,
|
|
on_close_callback,
|
|
nbio,
|
|
internal_connection_workflow=internal_connection_workflow)
|
|
|
|
@classmethod
|
|
def create_connection(cls,
|
|
connection_configs,
|
|
on_done,
|
|
custom_ioloop=None,
|
|
workflow=None):
|
|
"""Implement
|
|
:py:classmethod::`pika.adapters.BaseConnection.create_connection()`.
|
|
"""
|
|
custom_ioloop = (custom_ioloop or
|
|
_GeventSelectorIOLoop(gevent.get_hub()))
|
|
|
|
nbio = _GeventSelectorIOServicesAdapter(custom_ioloop)
|
|
|
|
def connection_factory(params):
|
|
"""Connection factory."""
|
|
if params is None:
|
|
raise ValueError('Expected pika.connection.Parameters '
|
|
'instance, but got None in params arg.')
|
|
return cls(parameters=params,
|
|
custom_ioloop=nbio,
|
|
internal_connection_workflow=False)
|
|
|
|
return cls._start_connection_workflow(
|
|
connection_configs=connection_configs,
|
|
connection_factory=connection_factory,
|
|
nbio=nbio,
|
|
workflow=workflow,
|
|
on_done=on_done)
|
|
|
|
|
|
class _TSafeCallbackQueue(object):
|
|
"""Dispatch callbacks from any thread to be executed in the main thread
|
|
efficiently with IO events.
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""
|
|
:param _GeventSelectorIOLoop loop: IO loop to add callbacks to.
|
|
"""
|
|
# Thread-safe, blocking queue.
|
|
self._queue = queue.Queue()
|
|
# PIPE to trigger an event when the queue is ready.
|
|
self._read_fd, self._write_fd = os.pipe()
|
|
# Lock around writes to the PIPE in case some platform/implementation
|
|
# requires this.
|
|
self._write_lock = threading.RLock()
|
|
|
|
@property
|
|
def fd(self):
|
|
"""The file-descriptor to register for READ events in the IO loop."""
|
|
return self._read_fd
|
|
|
|
def add_callback_threadsafe(self, callback):
|
|
"""Add an item to the queue from any thread. The configured handler
|
|
will be invoked with the item in the main thread.
|
|
|
|
:param item: Object to add to the queue.
|
|
"""
|
|
self._queue.put(callback)
|
|
with self._write_lock:
|
|
# The value written is not important.
|
|
os.write(self._write_fd, b'\xFF')
|
|
|
|
def run_next_callback(self):
|
|
"""Invoke the next callback from the queue.
|
|
|
|
MUST run in the main thread. If no callback was added to the queue,
|
|
this will block the IO loop.
|
|
|
|
Performs a blocking READ on the pipe so must only be called when the
|
|
pipe is ready for reading.
|
|
"""
|
|
try:
|
|
callback = self._queue.get_nowait()
|
|
except queue.Empty:
|
|
# Should never happen.
|
|
LOGGER.warning("Callback queue was empty.")
|
|
else:
|
|
# Read the byte from the pipe so the event doesn't re-fire.
|
|
os.read(self._read_fd, 1)
|
|
callback()
|
|
|
|
|
|
class _GeventSelectorIOLoop(AbstractSelectorIOLoop):
|
|
"""Implementation of `AbstractSelectorIOLoop` using the Gevent event loop.
|
|
|
|
Required by implementations of `SelectorIOServicesAdapter`.
|
|
"""
|
|
# Gevent's READ and WRITE masks are defined as 1 and 2 respectively. No
|
|
# ERROR mask is defined.
|
|
# See http://www.gevent.org/api/gevent.hub.html#gevent._interfaces.ILoop.io
|
|
READ = 1
|
|
WRITE = 2
|
|
ERROR = 0
|
|
|
|
def __init__(self, gevent_hub=None):
|
|
"""
|
|
:param gevent._interfaces.ILoop gevent_loop:
|
|
"""
|
|
self._hub = gevent_hub or gevent.get_hub()
|
|
self._io_watchers_by_fd = {}
|
|
# Used to start/stop the loop.
|
|
self._waiter = gevent.hub.Waiter()
|
|
|
|
# For adding callbacks from other threads. See `add_callback(..)`.
|
|
self._callback_queue = _TSafeCallbackQueue()
|
|
|
|
def run_callback_in_main_thread(fd, events):
|
|
"""Swallow the fd and events arguments."""
|
|
del fd
|
|
del events
|
|
self._callback_queue.run_next_callback()
|
|
|
|
self.add_handler(self._callback_queue.fd, run_callback_in_main_thread,
|
|
self.READ)
|
|
|
|
def close(self):
|
|
"""Release the loop's resources."""
|
|
self._hub.loop.destroy()
|
|
self._hub = None
|
|
|
|
def start(self):
|
|
"""Run the I/O loop. It will loop until requested to exit. See `stop()`.
|
|
"""
|
|
LOGGER.debug("Passing control to Gevent's IOLoop")
|
|
self._waiter.get() # Block until 'stop()' is called.
|
|
|
|
LOGGER.debug("Control was passed back from Gevent's IOLoop")
|
|
self._waiter.clear()
|
|
|
|
def stop(self):
|
|
"""Request exit from the ioloop. The loop is NOT guaranteed to
|
|
stop before this method returns.
|
|
|
|
To invoke `stop()` safely from a thread other than this IOLoop's thread,
|
|
call it via `add_callback_threadsafe`; e.g.,
|
|
|
|
`ioloop.add_callback(ioloop.stop)`
|
|
"""
|
|
self._waiter.switch(None)
|
|
|
|
def add_callback(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 in IOLoop. All other
|
|
manipulations of IOLoop must be performed from the IOLoop's thread.
|
|
|
|
For example, a thread may request a call to the `stop` method of an
|
|
ioloop that is running in a different thread via
|
|
`ioloop.add_callback_threadsafe(ioloop.stop)`
|
|
|
|
:param callable callback: The callback method
|
|
"""
|
|
if gevent.get_hub() == self._hub:
|
|
# We're in the main thread; just add the callback.
|
|
LOGGER.debug("Adding callback from main thread")
|
|
self._hub.loop.run_callback(callback)
|
|
else:
|
|
# This isn't the main thread and Gevent's hub/loop don't provide
|
|
# any thread-safety so enqueue the callback for it to be registered
|
|
# in the main thread.
|
|
LOGGER.debug("Adding callback from another thread")
|
|
callback = functools.partial(self._hub.loop.run_callback, callback)
|
|
self._callback_queue.add_callback_threadsafe(callback)
|
|
|
|
def call_later(self, delay, callback):
|
|
"""Add the callback to the IOLoop timer to be called after delay seconds
|
|
from the time of call on best-effort basis. Returns a handle to the
|
|
timeout.
|
|
|
|
:param float delay: The number of seconds to wait to call callback
|
|
:param callable callback: The callback method
|
|
:returns: handle to the created timeout that may be passed to
|
|
`remove_timeout()`
|
|
:rtype: object
|
|
"""
|
|
timer = self._hub.loop.timer(delay)
|
|
timer.start(callback)
|
|
return timer
|
|
|
|
def remove_timeout(self, timeout_handle):
|
|
"""Remove a timeout
|
|
|
|
:param timeout_handle: Handle of timeout to remove
|
|
"""
|
|
timeout_handle.close()
|
|
|
|
def add_handler(self, fd, handler, events):
|
|
"""Start watching the given file descriptor for events
|
|
|
|
:param int fd: The file descriptor
|
|
:param callable handler: When requested event(s) occur,
|
|
`handler(fd, events)` will be called.
|
|
:param int events: The event mask (READ|WRITE)
|
|
"""
|
|
io_watcher = self._hub.loop.io(fd, events)
|
|
self._io_watchers_by_fd[fd] = io_watcher
|
|
io_watcher.start(handler, fd, events)
|
|
|
|
def update_handler(self, fd, events):
|
|
"""Change the events being watched for.
|
|
|
|
:param int fd: The file descriptor
|
|
:param int events: The new event mask (READ|WRITE)
|
|
"""
|
|
io_watcher = self._io_watchers_by_fd[fd]
|
|
# Save callback from the original watcher. The close the old watcher
|
|
# and create a new one using the saved callback and the new events.
|
|
callback = io_watcher.callback
|
|
io_watcher.close()
|
|
del self._io_watchers_by_fd[fd]
|
|
self.add_handler(fd, callback, events)
|
|
|
|
def remove_handler(self, fd):
|
|
"""Stop watching the given file descriptor for events
|
|
|
|
:param int fd: The file descriptor
|
|
"""
|
|
io_watcher = self._io_watchers_by_fd[fd]
|
|
io_watcher.close()
|
|
del self._io_watchers_by_fd[fd]
|
|
|
|
|
|
class _GeventSelectorIOServicesAdapter(SelectorIOServicesAdapter):
|
|
"""SelectorIOServicesAdapter implementation using Gevent's DNS resolver."""
|
|
|
|
def getaddrinfo(self,
|
|
host,
|
|
port,
|
|
on_done,
|
|
family=0,
|
|
socktype=0,
|
|
proto=0,
|
|
flags=0):
|
|
"""Implement :py:meth:`.nbio_interface.AbstractIOServices.getaddrinfo()`.
|
|
"""
|
|
resolver = _GeventAddressResolver(native_loop=self._loop,
|
|
host=host,
|
|
port=port,
|
|
family=family,
|
|
socktype=socktype,
|
|
proto=proto,
|
|
flags=flags,
|
|
on_done=on_done)
|
|
resolver.start()
|
|
# Return needs an implementation of `AbstractIOReference`.
|
|
return _GeventIOLoopIOHandle(resolver)
|
|
|
|
|
|
class _GeventIOLoopIOHandle(AbstractIOReference):
|
|
"""Implement `AbstractIOReference`.
|
|
|
|
Only used to wrap the _GeventAddressResolver.
|
|
"""
|
|
|
|
def __init__(self, subject):
|
|
"""
|
|
:param subject: subject of the reference containing a `cancel()` method
|
|
"""
|
|
self._cancel = subject.cancel
|
|
|
|
def cancel(self):
|
|
"""Cancel pending operation
|
|
|
|
:returns: False if was already done or cancelled; True otherwise
|
|
:rtype: bool
|
|
"""
|
|
return self._cancel()
|
|
|
|
|
|
class _GeventAddressResolver(object):
|
|
"""Performs getaddrinfo asynchronously Gevent's configured resolver in a
|
|
separate greenlet and invoking the provided callback with the result.
|
|
|
|
See: http://www.gevent.org/dns.html
|
|
"""
|
|
__slots__ = (
|
|
'_loop',
|
|
'_on_done',
|
|
'_greenlet',
|
|
# getaddrinfo(..) args:
|
|
'_ga_host',
|
|
'_ga_port',
|
|
'_ga_family',
|
|
'_ga_socktype',
|
|
'_ga_proto',
|
|
'_ga_flags')
|
|
|
|
def __init__(self, native_loop, host, port, family, socktype, proto, flags,
|
|
on_done):
|
|
"""Initialize the `_GeventAddressResolver`.
|
|
|
|
:param AbstractSelectorIOLoop native_loop:
|
|
:param host: `see socket.getaddrinfo()`
|
|
:param port: `see socket.getaddrinfo()`
|
|
:param family: `see socket.getaddrinfo()`
|
|
:param socktype: `see socket.getaddrinfo()`
|
|
:param proto: `see socket.getaddrinfo()`
|
|
:param flags: `see socket.getaddrinfo()`
|
|
:param on_done: on_done(records|BaseException) callback for reporting
|
|
result from the given I/O loop. The single arg will be either an
|
|
exception object (check for `BaseException`) in case of failure or
|
|
the result returned by `socket.getaddrinfo()`.
|
|
"""
|
|
check_callback_arg(on_done, 'on_done')
|
|
|
|
self._loop = native_loop
|
|
self._on_done = on_done
|
|
# Reference to the greenlet performing `getaddrinfo`.
|
|
self._greenlet = None
|
|
# getaddrinfo(..) args.
|
|
self._ga_host = host
|
|
self._ga_port = port
|
|
self._ga_family = family
|
|
self._ga_socktype = socktype
|
|
self._ga_proto = proto
|
|
self._ga_flags = flags
|
|
|
|
def start(self):
|
|
"""Start an asynchronous getaddrinfo invocation."""
|
|
if self._greenlet is None:
|
|
self._greenlet = gevent.spawn_raw(self._resolve)
|
|
else:
|
|
LOGGER.warning("_GeventAddressResolver already started")
|
|
|
|
def cancel(self):
|
|
"""Cancel the pending resolver."""
|
|
changed = False
|
|
|
|
if self._greenlet is not None:
|
|
changed = True
|
|
self._stop_greenlet()
|
|
|
|
self._cleanup()
|
|
return changed
|
|
|
|
def _cleanup(self):
|
|
"""Stop the resolver and release any resources."""
|
|
self._stop_greenlet()
|
|
self._loop = None
|
|
self._on_done = None
|
|
|
|
def _stop_greenlet(self):
|
|
"""Stop the greenlet performing getaddrinfo if running.
|
|
|
|
Otherwise, this is a no-op.
|
|
"""
|
|
if self._greenlet is not None:
|
|
gevent.kill(self._greenlet)
|
|
self._greenlet = None
|
|
|
|
def _resolve(self):
|
|
"""Call `getaddrinfo()` and return result via user's callback
|
|
function on the configured IO loop.
|
|
"""
|
|
try:
|
|
# NOTE(JG): Can't use kwargs with getaddrinfo on Python <= v2.7.
|
|
result = gevent.socket.getaddrinfo(self._ga_host, self._ga_port,
|
|
self._ga_family,
|
|
self._ga_socktype,
|
|
self._ga_proto, self._ga_flags)
|
|
except Exception as exc: # pylint: disable=broad-except
|
|
LOGGER.error('Address resolution failed: %r', exc)
|
|
result = exc
|
|
|
|
callback = functools.partial(self._dispatch_callback, result)
|
|
self._loop.add_callback(callback)
|
|
|
|
def _dispatch_callback(self, result):
|
|
"""Invoke the configured completion callback and any subsequent cleanup.
|
|
|
|
:param result: result from getaddrinfo, or the exception if raised.
|
|
"""
|
|
try:
|
|
LOGGER.debug(
|
|
'Invoking async getaddrinfo() completion callback; host=%r',
|
|
self._ga_host)
|
|
self._on_done(result)
|
|
finally:
|
|
self._cleanup()
|