Gevent driven zmq poller, per sage advice of benoitc.

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)

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()

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)