PYTHON   90

core

Guest on 30th April 2022 11:54:15 PM

  1. #-----------------------------------------------------------------------------
  2. #  Copyright (C) Travis Cline
  3. #
  4. #  This file is part of pyzmq
  5. #  It is adapted from upstream project zeromq_gevent under the New BSD License
  6. #
  7. #  Distributed under the terms of the New BSD License.  The full license is in
  8. #  the file COPYING.BSD, distributed as part of this software.
  9. #-----------------------------------------------------------------------------
  10.  
  11. """This module wraps the :class:`Socket` and :class:`Context` found in :mod:`pyzmq <zmq>` to be non blocking
  12. """
  13.  
  14. from __future__ import print_function
  15.  
  16. import sys
  17. import time
  18. import warnings
  19.  
  20. import zmq
  21.  
  22. from zmq import Context as _original_Context
  23. from zmq import Socket as _original_Socket
  24. from .poll import _Poller
  25.  
  26. import gevent
  27. from gevent.event import AsyncResult
  28. from gevent.hub import get_hub
  29.  
  30. if hasattr(zmq, 'RCVTIMEO'):
  31.     TIMEOS = (zmq.RCVTIMEO, zmq.SNDTIMEO)
  32. else:
  33.     TIMEOS = ()
  34.  
  35. def _stop(evt):
  36.     """simple wrapper for stopping an Event, allowing for method rename in gevent 1.0"""
  37.     try:
  38.         evt.stop()
  39.     except AttributeError as e:
  40.         # gevent<1.0 compat
  41.         evt.cancel()
  42.  
  43. class _Socket(_original_Socket):
  44.     """Green version of :class:`zmq.Socket`
  45.  
  46.    The following methods are overridden:
  47.  
  48.        * send
  49.        * recv
  50.  
  51.    To ensure that the ``zmq.NOBLOCK`` flag is set and that sending or receiving
  52.    is deferred to the hub if a ``zmq.EAGAIN`` (retry) error is raised.
  53.    
  54.    The `__state_changed` method is triggered when the zmq.FD for the socket is
  55.    marked as readable and triggers the necessary read and write events (which
  56.    are waited for in the recv and send methods).
  57.  
  58.    Some double underscore prefixes are used to minimize pollution of
  59.    :class:`zmq.Socket`'s namespace.
  60.    """
  61.     __in_send_multipart = False
  62.     __in_recv_multipart = False
  63.     __writable = None
  64.     __readable = None
  65.     _state_event = None
  66.     _gevent_bug_timeout = 11.6 # timeout for not trusting gevent
  67.     _debug_gevent = False # turn on if you think gevent is missing events
  68.     _poller_class = _Poller
  69.    
  70.     def __init__(self, *a, **kw):
  71.         super(_Socket, self).__init__(*a, **kw)
  72.         self.__in_send_multipart = False
  73.         self.__in_recv_multipart = False
  74.         self.__setup_events()
  75.        
  76.  
  77.     def __del__(self):
  78.         self.close()
  79.  
  80.     def close(self, linger=None):
  81.         super(_Socket, self).close(linger)
  82.         self.__cleanup_events()
  83.  
  84.     def __cleanup_events(self):
  85.         # close the _state_event event, keeps the number of active file descriptors down
  86.         if getattr(self, '_state_event', None):
  87.             _stop(self._state_event)
  88.             self._state_event = None
  89.         # if the socket has entered a close state resume any waiting greenlets
  90.         self.__writable.set()
  91.         self.__readable.set()
  92.  
  93.     def __setup_events(self):
  94.         self.__readable = AsyncResult()
  95.         self.__writable = AsyncResult()
  96.         self.__readable.set()
  97.         self.__writable.set()
  98.        
  99.         try:
  100.             self._state_event = get_hub().loop.io(self.getsockopt(zmq.FD), 1) # read state watcher
  101.             self._state_event.start(self.__state_changed)
  102.         except AttributeError:
  103.             # for gevent<1.0 compatibility
  104.             from gevent.core import read_event
  105.             self._state_event = read_event(self.getsockopt(zmq.FD), self.__state_changed, persist=True)
  106.  
  107.     def __state_changed(self, event=None, _evtype=None):
  108.         if self.closed:
  109.             self.__cleanup_events()
  110.             return
  111.         try:
  112.             # avoid triggering __state_changed from inside __state_changed
  113.             events = super(_Socket, self).getsockopt(zmq.EVENTS)
  114.         except zmq.ZMQError as exc:
  115.             self.__writable.set_exception(exc)
  116.             self.__readable.set_exception(exc)
  117.         else:
  118.             if events & zmq.POLLOUT:
  119.                 self.__writable.set()
  120.             if events & zmq.POLLIN:
  121.                 self.__readable.set()
  122.  
  123.     def _wait_write(self):
  124.         assert self.__writable.ready(), "Only one greenlet can be waiting on this event"
  125.         self.__writable = AsyncResult()
  126.         # timeout is because libzmq cannot be trusted to properly signal a new send event:
  127.         # this is effectively a maximum poll interval of 1s
  128.         tic = time.time()
  129.         dt = self._gevent_bug_timeout
  130.         if dt:
  131.             timeout = gevent.Timeout(seconds=dt)
  132.         else:
  133.             timeout = None
  134.         try:
  135.             if timeout:
  136.                 timeout.start()
  137.             self.__writable.get(block=True)
  138.         except gevent.Timeout as t:
  139.             if t is not timeout:
  140.                 raise
  141.             toc = time.time()
  142.             # gevent bug: get can raise timeout even on clean return
  143.             # don't display zmq bug warning for gevent bug (this is getting ridiculous)
  144.             if self._debug_gevent and timeout and toc-tic > dt and \
  145.                     self.getsockopt(zmq.EVENTS) & zmq.POLLOUT:
  146.                 print("BUG: gevent may have missed a libzmq send event on %i!" % self.FD, file=sys.stderr)
  147.         finally:
  148.             if timeout:
  149.                 timeout.cancel()
  150.             self.__writable.set()
  151.  
  152.     def _wait_read(self):
  153.         assert self.__readable.ready(), "Only one greenlet can be waiting on this event"
  154.         self.__readable = AsyncResult()
  155.         # timeout is because libzmq cannot always be trusted to play nice with libevent.
  156.         # I can only confirm that this actually happens for send, but lets be symmetrical
  157.         # with our dirty hacks.
  158.         # this is effectively a maximum poll interval of 1s
  159.         tic = time.time()
  160.         dt = self._gevent_bug_timeout
  161.         if dt:
  162.             timeout = gevent.Timeout(seconds=dt)
  163.         else:
  164.             timeout = None
  165.         try:
  166.             if timeout:
  167.                 timeout.start()
  168.             self.__readable.get(block=True)
  169.         except gevent.Timeout as t:
  170.             if t is not timeout:
  171.                 raise
  172.             toc = time.time()
  173.             # gevent bug: get can raise timeout even on clean return
  174.             # don't display zmq bug warning for gevent bug (this is getting ridiculous)
  175.             if self._debug_gevent and timeout and toc-tic > dt and \
  176.                     self.getsockopt(zmq.EVENTS) & zmq.POLLIN:
  177.                 print("BUG: gevent may have missed a libzmq recv event on %i!" % self.FD, file=sys.stderr)
  178.         finally:
  179.             if timeout:
  180.                 timeout.cancel()
  181.             self.__readable.set()
  182.  
  183.     def send(self, data, flags=0, copy=True, track=False):
  184.         """send, which will only block current greenlet
  185.        
  186.        state_changed always fires exactly once (success or fail) at the
  187.        end of this method.
  188.        """
  189.        
  190.         # if we're given the NOBLOCK flag act as normal and let the EAGAIN get raised
  191.         if flags & zmq.NOBLOCK:
  192.             try:
  193.                 msg = super(_Socket, self).send(data, flags, copy, track)
  194.             finally:
  195.                 if not self.__in_send_multipart:
  196.                     self.__state_changed()
  197.             return msg
  198.         # ensure the zmq.NOBLOCK flag is part of flags
  199.         flags |= zmq.NOBLOCK
  200.         while True: # Attempt to complete this operation indefinitely, blocking the current greenlet
  201.             try:
  202.                 # attempt the actual call
  203.                 msg = super(_Socket, self).send(data, flags, copy, track)
  204.             except zmq.ZMQError as e:
  205.                 # if the raised ZMQError is not EAGAIN, reraise
  206.                 if e.errno != zmq.EAGAIN:
  207.                     if not self.__in_send_multipart:
  208.                         self.__state_changed()
  209.                     raise
  210.             else:
  211.                 if not self.__in_send_multipart:
  212.                     self.__state_changed()
  213.                 return msg
  214.             # defer to the event loop until we're notified the socket is writable
  215.             self._wait_write()
  216.  
  217.     def recv(self, flags=0, copy=True, track=False):
  218.         """recv, which will only block current greenlet
  219.        
  220.        state_changed always fires exactly once (success or fail) at the
  221.        end of this method.
  222.        """
  223.         if flags & zmq.NOBLOCK:
  224.             try:
  225.                 msg = super(_Socket, self).recv(flags, copy, track)
  226.             finally:
  227.                 if not self.__in_recv_multipart:
  228.                     self.__state_changed()
  229.             return msg
  230.        
  231.         flags |= zmq.NOBLOCK
  232.         while True:
  233.             try:
  234.                 msg = super(_Socket, self).recv(flags, copy, track)
  235.             except zmq.ZMQError as e:
  236.                 if e.errno != zmq.EAGAIN:
  237.                     if not self.__in_recv_multipart:
  238.                         self.__state_changed()
  239.                     raise
  240.             else:
  241.                 if not self.__in_recv_multipart:
  242.                     self.__state_changed()
  243.                 return msg
  244.             self._wait_read()
  245.    
  246.     def send_multipart(self, *args, **kwargs):
  247.         """wrap send_multipart to prevent state_changed on each partial send"""
  248.         self.__in_send_multipart = True
  249.         try:
  250.             msg = super(_Socket, self).send_multipart(*args, **kwargs)
  251.         finally:
  252.             self.__in_send_multipart = False
  253.             self.__state_changed()
  254.         return msg
  255.    
  256.     def recv_multipart(self, *args, **kwargs):
  257.         """wrap recv_multipart to prevent state_changed on each partial recv"""
  258.         self.__in_recv_multipart = True
  259.         try:
  260.             msg = super(_Socket, self).recv_multipart(*args, **kwargs)
  261.         finally:
  262.             self.__in_recv_multipart = False
  263.             self.__state_changed()
  264.         return msg
  265.    
  266.     def get(self, opt):
  267.         """trigger state_changed on getsockopt(EVENTS)"""
  268.         if opt in TIMEOS:
  269.             warnings.warn("TIMEO socket options have no effect in zmq.green", UserWarning)
  270.         optval = super(_Socket, self).get(opt)
  271.         if opt == zmq.EVENTS:
  272.             self.__state_changed()
  273.         return optval
  274.    
  275.     def set(self, opt, val):
  276.         """set socket option"""
  277.         if opt in TIMEOS:
  278.             warnings.warn("TIMEO socket options have no effect in zmq.green", UserWarning)
  279.         return super(_Socket, self).set(opt, val)
  280.  
  281.  
  282. class _Context(_original_Context):
  283.     """Replacement for :class:`zmq.Context`
  284.  
  285.    Ensures that the greened Socket above is used in calls to `socket`.
  286.    """
  287.     _socket_class = _Socket

Raw Paste


Login or Register to edit or fork this paste. It's free.