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