wait.py 5.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153
  1. import errno
  2. import select
  3. import sys
  4. from functools import partial
  5. try:
  6. from time import monotonic
  7. except ImportError:
  8. from time import time as monotonic
  9. __all__ = ["NoWayToWaitForSocketError", "wait_for_read", "wait_for_write"]
  10. class NoWayToWaitForSocketError(Exception):
  11. pass
  12. # How should we wait on sockets?
  13. #
  14. # There are two types of APIs you can use for waiting on sockets: the fancy
  15. # modern stateful APIs like epoll/kqueue, and the older stateless APIs like
  16. # select/poll. The stateful APIs are more efficient when you have a lots of
  17. # sockets to keep track of, because you can set them up once and then use them
  18. # lots of times. But we only ever want to wait on a single socket at a time
  19. # and don't want to keep track of state, so the stateless APIs are actually
  20. # more efficient. So we want to use select() or poll().
  21. #
  22. # Now, how do we choose between select() and poll()? On traditional Unixes,
  23. # select() has a strange calling convention that makes it slow, or fail
  24. # altogether, for high-numbered file descriptors. The point of poll() is to fix
  25. # that, so on Unixes, we prefer poll().
  26. #
  27. # On Windows, there is no poll() (or at least Python doesn't provide a wrapper
  28. # for it), but that's OK, because on Windows, select() doesn't have this
  29. # strange calling convention; plain select() works fine.
  30. #
  31. # So: on Windows we use select(), and everywhere else we use poll(). We also
  32. # fall back to select() in case poll() is somehow broken or missing.
  33. if sys.version_info >= (3, 5):
  34. # Modern Python, that retries syscalls by default
  35. def _retry_on_intr(fn, timeout):
  36. return fn(timeout)
  37. else:
  38. # Old and broken Pythons.
  39. def _retry_on_intr(fn, timeout):
  40. if timeout is None:
  41. deadline = float("inf")
  42. else:
  43. deadline = monotonic() + timeout
  44. while True:
  45. try:
  46. return fn(timeout)
  47. # OSError for 3 <= pyver < 3.5, select.error for pyver <= 2.7
  48. except (OSError, select.error) as e:
  49. # 'e.args[0]' incantation works for both OSError and select.error
  50. if e.args[0] != errno.EINTR:
  51. raise
  52. else:
  53. timeout = deadline - monotonic()
  54. if timeout < 0:
  55. timeout = 0
  56. if timeout == float("inf"):
  57. timeout = None
  58. continue
  59. def select_wait_for_socket(sock, read=False, write=False, timeout=None):
  60. if not read and not write:
  61. raise RuntimeError("must specify at least one of read=True, write=True")
  62. rcheck = []
  63. wcheck = []
  64. if read:
  65. rcheck.append(sock)
  66. if write:
  67. wcheck.append(sock)
  68. # When doing a non-blocking connect, most systems signal success by
  69. # marking the socket writable. Windows, though, signals success by marked
  70. # it as "exceptional". We paper over the difference by checking the write
  71. # sockets for both conditions. (The stdlib selectors module does the same
  72. # thing.)
  73. fn = partial(select.select, rcheck, wcheck, wcheck)
  74. rready, wready, xready = _retry_on_intr(fn, timeout)
  75. return bool(rready or wready or xready)
  76. def poll_wait_for_socket(sock, read=False, write=False, timeout=None):
  77. if not read and not write:
  78. raise RuntimeError("must specify at least one of read=True, write=True")
  79. mask = 0
  80. if read:
  81. mask |= select.POLLIN
  82. if write:
  83. mask |= select.POLLOUT
  84. poll_obj = select.poll()
  85. poll_obj.register(sock, mask)
  86. # For some reason, poll() takes timeout in milliseconds
  87. def do_poll(t):
  88. if t is not None:
  89. t *= 1000
  90. return poll_obj.poll(t)
  91. return bool(_retry_on_intr(do_poll, timeout))
  92. def null_wait_for_socket(*args, **kwargs):
  93. raise NoWayToWaitForSocketError("no select-equivalent available")
  94. def _have_working_poll():
  95. # Apparently some systems have a select.poll that fails as soon as you try
  96. # to use it, either due to strange configuration or broken monkeypatching
  97. # from libraries like eventlet/greenlet.
  98. try:
  99. poll_obj = select.poll()
  100. _retry_on_intr(poll_obj.poll, 0)
  101. except (AttributeError, OSError):
  102. return False
  103. else:
  104. return True
  105. def wait_for_socket(*args, **kwargs):
  106. # We delay choosing which implementation to use until the first time we're
  107. # called. We could do it at import time, but then we might make the wrong
  108. # decision if someone goes wild with monkeypatching select.poll after
  109. # we're imported.
  110. global wait_for_socket
  111. if _have_working_poll():
  112. wait_for_socket = poll_wait_for_socket
  113. elif hasattr(select, "select"):
  114. wait_for_socket = select_wait_for_socket
  115. else: # Platform-specific: Appengine.
  116. wait_for_socket = null_wait_for_socket
  117. return wait_for_socket(*args, **kwargs)
  118. def wait_for_read(sock, timeout=None):
  119. """Waits for reading to be available on a given socket.
  120. Returns True if the socket is readable, or False if the timeout expired.
  121. """
  122. return wait_for_socket(sock, read=True, timeout=timeout)
  123. def wait_for_write(sock, timeout=None):
  124. """Waits for writing to be available on a given socket.
  125. Returns True if the socket is readable, or False if the timeout expired.
  126. """
  127. return wait_for_socket(sock, write=True, timeout=timeout)