diff --git a/zipline/component.py b/zipline/component.py index d33ae44b..cfdc8781 100644 --- a/zipline/component.py +++ b/zipline/component.py @@ -22,9 +22,11 @@ import gevent_zeromq from datetime import datetime import zipline.util as qutil +from zipline.gpoll import _Poller as GeventPoller from zipline.protocol import CONTROL_PROTOCOL, COMPONENT_STATE, \ COMPONENT_FAILURE, BACKTEST_STATE + class Component(object): """ Base class for components. Defines the the base messaging @@ -148,21 +150,24 @@ class Component(object): if flavor == 'mp': self.zmq = zmq self.context = self.zmq.Context() + self.zmq_poller = self.zmq.Poller return if flavor == 'thread': self.zmq = zmq self.context = self.zmq.Context.instance() + self.zmq_poller = self.zmq.Poller return if flavor == 'green': self.zmq = gevent_zeromq.zmq self.context = self.zmq.Context.instance() + self.zmq_poller = GeventPoller return if flavor == 'pypy': self.zmq = zmq self.context = self.zmq.Context.instance() + self.zmq_poller = self.zmq.Poller return - import pdb; pdb.set_trace() raise Exception("Unknown ZeroMQ Flavor") def _run(self): @@ -335,7 +340,9 @@ class Component(object): handling sockets. """ - self.poll = self.zmq.Poller() + # Initializes the poller class specified by the flavor of + # ZeroMQ. Either zmq.Poller or gpoll.Poller . + self.poll = self.zmq_poller() def receive_sync_ack(self): """ @@ -437,11 +444,7 @@ class Component(object): self.sync_socket.connect(self.addresses['sync_address']) #self.sync_socket.setsockopt(self.zmq.LINGER,0) - # Explictly a different poller for obvious reasons. - # I'm not fond of having this poller init'd as a side - # effect of a method call. Still thinking about where to - # put it at the moment though... - self.sync_poller = self.zmq.Poller() + self.sync_poller = self.zmq_poller() self.sync_poller.register(self.sync_socket, self.zmq.POLLIN) self.sockets.append(self.sync_socket) diff --git a/zipline/gpoll.py b/zipline/gpoll.py new file mode 100644 index 00000000..8c18d4cf --- /dev/null +++ b/zipline/gpoll.py @@ -0,0 +1,99 @@ +""" +This is somewhat legally ambigious, since it technically +hasn't been merged in gevent_zeromq but given that the +author issued it as a Pull Request on a MIT project, +indicates that its probably fine to use. ~Steve +""" + +import zmq +from zmq import * + +from zmq.core.poll import Poller as _original_Poller + +import gevent +from gevent import select +from gevent_zeromq.core import _Socket + +def patch_poller(self): + zmq.Poller = _Poller + +class _Poller(_original_Poller): + """ + Replacement for :class:`zmq.core.Poller` + + Ensures that the greened Poller below is used in calls + to :meth:`zmq.core.Poller.poll`. + """ + + def _get_descriptors(self): + """ + Returns three elements tuple with socket descriptors ready for + gevent.select + """ + rlist = [] + wlist = [] + xlist = [] + + for socket, flags in self.sockets.items(): + if isinstance(socket, _Socket): + fd = socket.getsockopt(FD) + 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 & POLLIN: rlist.append(fd) + if flags & POLLOUT: wlist.append(fd) + if flags & 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 + + timeout = int(timeout) + 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 + while True: + events = super(_Poller, self).poll(0) + if events or timeout == 0: + return events + + # wait for activity on sockets in a green way + if not rlist and not wlist and not xlist: + rlist, wlist, xlist = self._get_descriptors() + + try: + select.select(rlist, wlist, xlist) + except gevent.select.error, ex: + raise ZMQError(*ex.args) + + except gevent.Timeout, t: + if t is not tout: + raise + return [] + finally: + if timeout > 0: + tout.cancel() diff --git a/zipline/messaging.py b/zipline/messaging.py index 1b7ffbdf..5e6ff570 100644 --- a/zipline/messaging.py +++ b/zipline/messaging.py @@ -101,7 +101,7 @@ class ComponentHost(Component): self.sync_socket = self.context.socket(self.zmq.REP) self.sync_socket.bind(self.addresses['sync_address']) - self.sync_poller = self.zmq.Poller() + self.sync_poller = self.zmq_poller() self.sync_poller.register(self.sync_socket, self.zmq.POLLIN) self.sockets.append(self.sync_socket)