98 lines
2.9 KiB
Python
98 lines
2.9 KiB
Python
|
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)
|
||
|
|
else:
|
||
|
|
tout = None
|
||
|
|
|
||
|
|
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()
|
||
|
|
|