123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153 |
- import errno
- from functools import partial
- import select
- import sys
- try:
- from time import monotonic
- except ImportError:
- from time import time as monotonic
- __all__ = ["NoWayToWaitForSocketError", "wait_for_read", "wait_for_write"]
- class NoWayToWaitForSocketError(Exception):
- pass
- # How should we wait on sockets?
- #
- # There are two types of APIs you can use for waiting on sockets: the fancy
- # modern stateful APIs like epoll/kqueue, and the older stateless APIs like
- # select/poll. The stateful APIs are more efficient when you have a lots of
- # sockets to keep track of, because you can set them up once and then use them
- # lots of times. But we only ever want to wait on a single socket at a time
- # and don't want to keep track of state, so the stateless APIs are actually
- # more efficient. So we want to use select() or poll().
- #
- # Now, how do we choose between select() and poll()? On traditional Unixes,
- # select() has a strange calling convention that makes it slow, or fail
- # altogether, for high-numbered file descriptors. The point of poll() is to fix
- # that, so on Unixes, we prefer poll().
- #
- # On Windows, there is no poll() (or at least Python doesn't provide a wrapper
- # for it), but that's OK, because on Windows, select() doesn't have this
- # strange calling convention; plain select() works fine.
- #
- # So: on Windows we use select(), and everywhere else we use poll(). We also
- # fall back to select() in case poll() is somehow broken or missing.
- if sys.version_info >= (3, 5):
- # Modern Python, that retries syscalls by default
- def _retry_on_intr(fn, timeout):
- return fn(timeout)
- else:
- # Old and broken Pythons.
- def _retry_on_intr(fn, timeout):
- if timeout is None:
- deadline = float("inf")
- else:
- deadline = monotonic() + timeout
- while True:
- try:
- return fn(timeout)
- # OSError for 3 <= pyver < 3.5, select.error for pyver <= 2.7
- except (OSError, select.error) as e:
- # 'e.args[0]' incantation works for both OSError and select.error
- if e.args[0] != errno.EINTR:
- raise
- else:
- timeout = deadline - monotonic()
- if timeout < 0:
- timeout = 0
- if timeout == float("inf"):
- timeout = None
- continue
- def select_wait_for_socket(sock, read=False, write=False, timeout=None):
- if not read and not write:
- raise RuntimeError("must specify at least one of read=True, write=True")
- rcheck = []
- wcheck = []
- if read:
- rcheck.append(sock)
- if write:
- wcheck.append(sock)
- # When doing a non-blocking connect, most systems signal success by
- # marking the socket writable. Windows, though, signals success by marked
- # it as "exceptional". We paper over the difference by checking the write
- # sockets for both conditions. (The stdlib selectors module does the same
- # thing.)
- fn = partial(select.select, rcheck, wcheck, wcheck)
- rready, wready, xready = _retry_on_intr(fn, timeout)
- return bool(rready or wready or xready)
- def poll_wait_for_socket(sock, read=False, write=False, timeout=None):
- if not read and not write:
- raise RuntimeError("must specify at least one of read=True, write=True")
- mask = 0
- if read:
- mask |= select.POLLIN
- if write:
- mask |= select.POLLOUT
- poll_obj = select.poll()
- poll_obj.register(sock, mask)
- # For some reason, poll() takes timeout in milliseconds
- def do_poll(t):
- if t is not None:
- t *= 1000
- return poll_obj.poll(t)
- return bool(_retry_on_intr(do_poll, timeout))
- def null_wait_for_socket(*args, **kwargs):
- raise NoWayToWaitForSocketError("no select-equivalent available")
- def _have_working_poll():
- # Apparently some systems have a select.poll that fails as soon as you try
- # to use it, either due to strange configuration or broken monkeypatching
- # from libraries like eventlet/greenlet.
- try:
- poll_obj = select.poll()
- _retry_on_intr(poll_obj.poll, 0)
- except (AttributeError, OSError):
- return False
- else:
- return True
- def wait_for_socket(*args, **kwargs):
- # We delay choosing which implementation to use until the first time we're
- # called. We could do it at import time, but then we might make the wrong
- # decision if someone goes wild with monkeypatching select.poll after
- # we're imported.
- global wait_for_socket
- if _have_working_poll():
- wait_for_socket = poll_wait_for_socket
- elif hasattr(select, "select"):
- wait_for_socket = select_wait_for_socket
- else: # Platform-specific: Appengine.
- wait_for_socket = null_wait_for_socket
- return wait_for_socket(*args, **kwargs)
- def wait_for_read(sock, timeout=None):
- """ Waits for reading to be available on a given socket.
- Returns True if the socket is readable, or False if the timeout expired.
- """
- return wait_for_socket(sock, read=True, timeout=timeout)
- def wait_for_write(sock, timeout=None):
- """ Waits for writing to be available on a given socket.
- Returns True if the socket is readable, or False if the timeout expired.
- """
- return wait_for_socket(sock, write=True, timeout=timeout)
|