diff options
Diffstat (limited to 'src/console/zmq/green')
| -rwxr-xr-x | src/console/zmq/green/__init__.py | 40 | ||||
| -rwxr-xr-x | src/console/zmq/green/core.py | 287 | ||||
| -rwxr-xr-x | src/console/zmq/green/device.py | 32 | ||||
| -rwxr-xr-x | src/console/zmq/green/eventloop/__init__.py | 3 | ||||
| -rwxr-xr-x | src/console/zmq/green/eventloop/ioloop.py | 33 | ||||
| -rwxr-xr-x | src/console/zmq/green/eventloop/zmqstream.py | 11 | ||||
| -rwxr-xr-x | src/console/zmq/green/poll.py | 95 |
7 files changed, 501 insertions, 0 deletions
diff --git a/src/console/zmq/green/__init__.py b/src/console/zmq/green/__init__.py new file mode 100755 index 00000000..ff7e5965 --- /dev/null +++ b/src/console/zmq/green/__init__.py @@ -0,0 +1,40 @@ +# -*- coding: utf-8 -*- +#----------------------------------------------------------------------------- +# Copyright (C) 2011-2012 Travis Cline +# +# This file is part of pyzmq +# It is adapted from upstream project zeromq_gevent under the New BSD License +# +# Distributed under the terms of the New BSD License. The full license is in +# the file COPYING.BSD, distributed as part of this software. +#----------------------------------------------------------------------------- + +"""zmq.green - gevent compatibility with zeromq. + +Usage +----- + +Instead of importing zmq directly, do so in the following manner: + +.. + + import zmq.green as zmq + + +Any calls that would have blocked the current thread will now only block the +current green thread. + +This compatibility is accomplished by ensuring the nonblocking flag is set +before any blocking operation and the ØMQ file descriptor is polled internally +to trigger needed events. +""" + +from zmq import * +from zmq.green.core import _Context, _Socket +from zmq.green.poll import _Poller +Context = _Context +Socket = _Socket +Poller = _Poller + +from zmq.green.device import device + diff --git a/src/console/zmq/green/core.py b/src/console/zmq/green/core.py new file mode 100755 index 00000000..9fc73e32 --- /dev/null +++ b/src/console/zmq/green/core.py @@ -0,0 +1,287 @@ +#----------------------------------------------------------------------------- +# Copyright (C) 2011-2012 Travis Cline +# +# This file is part of pyzmq +# It is adapted from upstream project zeromq_gevent under the New BSD License +# +# Distributed under the terms of the New BSD License. The full license is in +# the file COPYING.BSD, distributed as part of this software. +#----------------------------------------------------------------------------- + +"""This module wraps the :class:`Socket` and :class:`Context` found in :mod:`pyzmq <zmq>` to be non blocking +""" + +from __future__ import print_function + +import sys +import time +import warnings + +import zmq + +from zmq import Context as _original_Context +from zmq import Socket as _original_Socket +from .poll import _Poller + +import gevent +from gevent.event import AsyncResult +from gevent.hub import get_hub + +if hasattr(zmq, 'RCVTIMEO'): + TIMEOS = (zmq.RCVTIMEO, zmq.SNDTIMEO) +else: + TIMEOS = () + +def _stop(evt): + """simple wrapper for stopping an Event, allowing for method rename in gevent 1.0""" + try: + evt.stop() + except AttributeError as e: + # gevent<1.0 compat + evt.cancel() + +class _Socket(_original_Socket): + """Green version of :class:`zmq.Socket` + + The following methods are overridden: + + * send + * recv + + To ensure that the ``zmq.NOBLOCK`` flag is set and that sending or receiving + is deferred to the hub if a ``zmq.EAGAIN`` (retry) error is raised. + + The `__state_changed` method is triggered when the zmq.FD for the socket is + marked as readable and triggers the necessary read and write events (which + are waited for in the recv and send methods). + + Some double underscore prefixes are used to minimize pollution of + :class:`zmq.Socket`'s namespace. + """ + __in_send_multipart = False + __in_recv_multipart = False + __writable = None + __readable = None + _state_event = None + _gevent_bug_timeout = 11.6 # timeout for not trusting gevent + _debug_gevent = False # turn on if you think gevent is missing events + _poller_class = _Poller + + def __init__(self, context, socket_type): + _original_Socket.__init__(self, context, socket_type) + self.__in_send_multipart = False + self.__in_recv_multipart = False + self.__setup_events() + + + def __del__(self): + self.close() + + def close(self, linger=None): + super(_Socket, self).close(linger) + self.__cleanup_events() + + def __cleanup_events(self): + # close the _state_event event, keeps the number of active file descriptors down + if getattr(self, '_state_event', None): + _stop(self._state_event) + self._state_event = None + # if the socket has entered a close state resume any waiting greenlets + self.__writable.set() + self.__readable.set() + + def __setup_events(self): + self.__readable = AsyncResult() + self.__writable = AsyncResult() + self.__readable.set() + self.__writable.set() + + try: + self._state_event = get_hub().loop.io(self.getsockopt(zmq.FD), 1) # read state watcher + self._state_event.start(self.__state_changed) + except AttributeError: + # for gevent<1.0 compatibility + from gevent.core import read_event + self._state_event = read_event(self.getsockopt(zmq.FD), self.__state_changed, persist=True) + + def __state_changed(self, event=None, _evtype=None): + if self.closed: + self.__cleanup_events() + return + try: + # avoid triggering __state_changed from inside __state_changed + events = super(_Socket, self).getsockopt(zmq.EVENTS) + except zmq.ZMQError as exc: + self.__writable.set_exception(exc) + self.__readable.set_exception(exc) + else: + if events & zmq.POLLOUT: + self.__writable.set() + if events & zmq.POLLIN: + self.__readable.set() + + def _wait_write(self): + assert self.__writable.ready(), "Only one greenlet can be waiting on this event" + self.__writable = AsyncResult() + # timeout is because libzmq cannot be trusted to properly signal a new send event: + # this is effectively a maximum poll interval of 1s + tic = time.time() + dt = self._gevent_bug_timeout + if dt: + timeout = gevent.Timeout(seconds=dt) + else: + timeout = None + try: + if timeout: + timeout.start() + self.__writable.get(block=True) + except gevent.Timeout as t: + if t is not timeout: + raise + toc = time.time() + # gevent bug: get can raise timeout even on clean return + # don't display zmq bug warning for gevent bug (this is getting ridiculous) + if self._debug_gevent and timeout and toc-tic > dt and \ + self.getsockopt(zmq.EVENTS) & zmq.POLLOUT: + print("BUG: gevent may have missed a libzmq send event on %i!" % self.FD, file=sys.stderr) + finally: + if timeout: + timeout.cancel() + self.__writable.set() + + def _wait_read(self): + assert self.__readable.ready(), "Only one greenlet can be waiting on this event" + self.__readable = AsyncResult() + # timeout is because libzmq cannot always be trusted to play nice with libevent. + # I can only confirm that this actually happens for send, but lets be symmetrical + # with our dirty hacks. + # this is effectively a maximum poll interval of 1s + tic = time.time() + dt = self._gevent_bug_timeout + if dt: + timeout = gevent.Timeout(seconds=dt) + else: + timeout = None + try: + if timeout: + timeout.start() + self.__readable.get(block=True) + except gevent.Timeout as t: + if t is not timeout: + raise + toc = time.time() + # gevent bug: get can raise timeout even on clean return + # don't display zmq bug warning for gevent bug (this is getting ridiculous) + if self._debug_gevent and timeout and toc-tic > dt and \ + self.getsockopt(zmq.EVENTS) & zmq.POLLIN: + print("BUG: gevent may have missed a libzmq recv event on %i!" % self.FD, file=sys.stderr) + finally: + if timeout: + timeout.cancel() + self.__readable.set() + + def send(self, data, flags=0, copy=True, track=False): + """send, which will only block current greenlet + + state_changed always fires exactly once (success or fail) at the + end of this method. + """ + + # if we're given the NOBLOCK flag act as normal and let the EAGAIN get raised + if flags & zmq.NOBLOCK: + try: + msg = super(_Socket, self).send(data, flags, copy, track) + finally: + if not self.__in_send_multipart: + self.__state_changed() + return msg + # ensure the zmq.NOBLOCK flag is part of flags + flags |= zmq.NOBLOCK + while True: # Attempt to complete this operation indefinitely, blocking the current greenlet + try: + # attempt the actual call + msg = super(_Socket, self).send(data, flags, copy, track) + except zmq.ZMQError as e: + # if the raised ZMQError is not EAGAIN, reraise + if e.errno != zmq.EAGAIN: + if not self.__in_send_multipart: + self.__state_changed() + raise + else: + if not self.__in_send_multipart: + self.__state_changed() + return msg + # defer to the event loop until we're notified the socket is writable + self._wait_write() + + def recv(self, flags=0, copy=True, track=False): + """recv, which will only block current greenlet + + state_changed always fires exactly once (success or fail) at the + end of this method. + """ + if flags & zmq.NOBLOCK: + try: + msg = super(_Socket, self).recv(flags, copy, track) + finally: + if not self.__in_recv_multipart: + self.__state_changed() + return msg + + flags |= zmq.NOBLOCK + while True: + try: + msg = super(_Socket, self).recv(flags, copy, track) + except zmq.ZMQError as e: + if e.errno != zmq.EAGAIN: + if not self.__in_recv_multipart: + self.__state_changed() + raise + else: + if not self.__in_recv_multipart: + self.__state_changed() + return msg + self._wait_read() + + def send_multipart(self, *args, **kwargs): + """wrap send_multipart to prevent state_changed on each partial send""" + self.__in_send_multipart = True + try: + msg = super(_Socket, self).send_multipart(*args, **kwargs) + finally: + self.__in_send_multipart = False + self.__state_changed() + return msg + + def recv_multipart(self, *args, **kwargs): + """wrap recv_multipart to prevent state_changed on each partial recv""" + self.__in_recv_multipart = True + try: + msg = super(_Socket, self).recv_multipart(*args, **kwargs) + finally: + self.__in_recv_multipart = False + self.__state_changed() + return msg + + def get(self, opt): + """trigger state_changed on getsockopt(EVENTS)""" + if opt in TIMEOS: + warnings.warn("TIMEO socket options have no effect in zmq.green", UserWarning) + optval = super(_Socket, self).get(opt) + if opt == zmq.EVENTS: + self.__state_changed() + return optval + + def set(self, opt, val): + """set socket option""" + if opt in TIMEOS: + warnings.warn("TIMEO socket options have no effect in zmq.green", UserWarning) + return super(_Socket, self).set(opt, val) + + +class _Context(_original_Context): + """Replacement for :class:`zmq.Context` + + Ensures that the greened Socket above is used in calls to `socket`. + """ + _socket_class = _Socket diff --git a/src/console/zmq/green/device.py b/src/console/zmq/green/device.py new file mode 100755 index 00000000..4b070237 --- /dev/null +++ b/src/console/zmq/green/device.py @@ -0,0 +1,32 @@ +# Copyright (C) PyZMQ Developers +# Distributed under the terms of the Modified BSD License. + +import zmq +from zmq.green import Poller + +def device(device_type, isocket, osocket): + """Start a zeromq device (gevent-compatible). + + Unlike the true zmq.device, this does not release the GIL. + + Parameters + ---------- + device_type : (QUEUE, FORWARDER, STREAMER) + The type of device to start (ignored). + isocket : Socket + The Socket instance for the incoming traffic. + osocket : Socket + The Socket instance for the outbound traffic. + """ + p = Poller() + if osocket == -1: + osocket = isocket + p.register(isocket, zmq.POLLIN) + p.register(osocket, zmq.POLLIN) + + while True: + events = dict(p.poll()) + if isocket in events: + osocket.send_multipart(isocket.recv_multipart()) + if osocket in events: + isocket.send_multipart(osocket.recv_multipart()) diff --git a/src/console/zmq/green/eventloop/__init__.py b/src/console/zmq/green/eventloop/__init__.py new file mode 100755 index 00000000..c5150efe --- /dev/null +++ b/src/console/zmq/green/eventloop/__init__.py @@ -0,0 +1,3 @@ +from zmq.green.eventloop.ioloop import IOLoop + +__all__ = ['IOLoop']
\ No newline at end of file diff --git a/src/console/zmq/green/eventloop/ioloop.py b/src/console/zmq/green/eventloop/ioloop.py new file mode 100755 index 00000000..e12fd5e9 --- /dev/null +++ b/src/console/zmq/green/eventloop/ioloop.py @@ -0,0 +1,33 @@ +from zmq.eventloop.ioloop import * +from zmq.green import Poller + +RealIOLoop = IOLoop +RealZMQPoller = ZMQPoller + +class IOLoop(RealIOLoop): + + def initialize(self, impl=None): + impl = _poll() if impl is None else impl + super(IOLoop, self).initialize(impl) + + @staticmethod + def instance(): + """Returns a global `IOLoop` instance. + + Most applications have a single, global `IOLoop` running on the + main thread. Use this method to get this instance from + another thread. To get the current thread's `IOLoop`, use `current()`. + """ + # install this class as the active IOLoop implementation + # when using tornado 3 + if tornado_version >= (3,): + PollIOLoop.configure(IOLoop) + return PollIOLoop.instance() + + +class ZMQPoller(RealZMQPoller): + """gevent-compatible version of ioloop.ZMQPoller""" + def __init__(self): + self._poller = Poller() + +_poll = ZMQPoller diff --git a/src/console/zmq/green/eventloop/zmqstream.py b/src/console/zmq/green/eventloop/zmqstream.py new file mode 100755 index 00000000..90fbd1f5 --- /dev/null +++ b/src/console/zmq/green/eventloop/zmqstream.py @@ -0,0 +1,11 @@ +from zmq.eventloop.zmqstream import * + +from zmq.green.eventloop.ioloop import IOLoop + +RealZMQStream = ZMQStream + +class ZMQStream(RealZMQStream): + + def __init__(self, socket, io_loop=None): + io_loop = io_loop or IOLoop.instance() + super(ZMQStream, self).__init__(socket, io_loop=io_loop) diff --git a/src/console/zmq/green/poll.py b/src/console/zmq/green/poll.py new file mode 100755 index 00000000..8f016129 --- /dev/null +++ b/src/console/zmq/green/poll.py @@ -0,0 +1,95 @@ +import zmq +import gevent +from gevent import select + +from zmq import Poller as _original_Poller + + +class _Poller(_original_Poller): + """Replacement for :class:`zmq.Poller` + + Ensures that the greened Poller below is used in calls to + :meth:`zmq.Poller.poll`. + """ + _gevent_bug_timeout = 1.33 # minimum poll interval, for working around gevent bug + + def _get_descriptors(self): + """Returns three elements tuple with socket descriptors ready + for gevent.select.select + """ + rlist = [] + wlist = [] + xlist = [] + + for socket, flags in self.sockets: + if isinstance(socket, zmq.Socket): + rlist.append(socket.getsockopt(zmq.FD)) + continue + elif isinstance(socket, int): + fd = socket + elif hasattr(socket, 'fileno'): + try: + fd = int(socket.fileno()) + except: + raise ValueError('fileno() must return an valid integer fd') + else: + raise TypeError('Socket must be a 0MQ socket, an integer fd ' + 'or have a fileno() method: %r' % socket) + + if flags & zmq.POLLIN: + rlist.append(fd) + if flags & zmq.POLLOUT: + wlist.append(fd) + if flags & zmq.POLLERR: + xlist.append(fd) + + return (rlist, wlist, xlist) + + def poll(self, timeout=-1): + """Overridden method to ensure that the green version of + Poller is used. + + Behaves the same as :meth:`zmq.core.Poller.poll` + """ + + if timeout is None: + timeout = -1 + + if timeout < 0: + timeout = -1 + + rlist = None + wlist = None + xlist = None + + if timeout > 0: + tout = gevent.Timeout.start_new(timeout/1000.0) + + try: + # Loop until timeout or events available + rlist, wlist, xlist = self._get_descriptors() + while True: + events = super(_Poller, self).poll(0) + if events or timeout == 0: + return events + + # wait for activity on sockets in a green way + # set a minimum poll frequency, + # because gevent < 1.0 cannot be trusted to catch edge-triggered FD events + _bug_timeout = gevent.Timeout.start_new(self._gevent_bug_timeout) + try: + select.select(rlist, wlist, xlist) + except gevent.Timeout as t: + if t is not _bug_timeout: + raise + finally: + _bug_timeout.cancel() + + except gevent.Timeout as t: + if t is not tout: + raise + return [] + finally: + if timeout > 0: + tout.cancel() + |