Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import errno
2import io
3import socket
4from io import SEEK_END
5from typing import Optional, Union
7from ..exceptions import ConnectionError, TimeoutError
8from ..utils import SENTINEL, SSL_AVAILABLE
10NONBLOCKING_EXCEPTION_ERROR_NUMBERS = {BlockingIOError: errno.EWOULDBLOCK}
12if SSL_AVAILABLE:
13 import ssl
15 if hasattr(ssl, "SSLWantReadError"):
16 NONBLOCKING_EXCEPTION_ERROR_NUMBERS[ssl.SSLWantReadError] = 2
17 NONBLOCKING_EXCEPTION_ERROR_NUMBERS[ssl.SSLWantWriteError] = 2
18 else:
19 NONBLOCKING_EXCEPTION_ERROR_NUMBERS[ssl.SSLError] = 2
21NONBLOCKING_EXCEPTIONS = tuple(NONBLOCKING_EXCEPTION_ERROR_NUMBERS.keys())
23SERVER_CLOSED_CONNECTION_ERROR = "Connection closed by server."
25SYM_CRLF = b"\r\n"
28class SocketBuffer:
29 def __init__(
30 self, socket: socket.socket, socket_read_size: int, socket_timeout: float
31 ):
32 self._sock = socket
33 self.socket_read_size = socket_read_size
34 self.socket_timeout = socket_timeout
35 self._buffer = io.BytesIO()
37 def _live_buffer(self) -> io.BytesIO:
38 """
39 The read buffer, or a ConnectionError when the connection is already gone.
41 ``close()`` closes the buffer before dropping it, and it can run while another
42 thread reads: the multi-database client closes connections from its health
43 check thread when a database is taken out of service, so a reader finds either
44 a closed ``BytesIO`` or ``None``. The connection is gone either way, so report
45 it the way every other teardown here is reported and let the retry layers act
46 on it, instead of surfacing ``ValueError: I/O operation on closed file`` to the
47 caller of the command.
48 """
49 buffer = self._buffer
51 if buffer is None or buffer.closed:
52 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
54 return buffer
56 def unread_bytes(self) -> int:
57 """
58 Remaining unread length of buffer
59 """
60 buffer = self._live_buffer()
62 try:
63 pos = buffer.tell()
64 end = buffer.seek(0, SEEK_END)
65 buffer.seek(pos)
66 except ValueError:
67 # Closed between the check above and here.
68 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
70 return end - pos
72 def _read_from_socket(
73 self,
74 length: Optional[int] = None,
75 timeout: Union[float, object] = SENTINEL,
76 raise_on_timeout: Optional[bool] = True,
77 ) -> bool:
78 sock = self._sock
79 socket_read_size = self.socket_read_size
80 marker = 0
81 custom_timeout = timeout is not SENTINEL
83 buf = self._live_buffer()
85 if sock is None:
86 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
88 try:
89 current_pos = buf.tell()
90 buf.seek(0, SEEK_END)
91 except ValueError:
92 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
94 if custom_timeout:
95 sock.settimeout(timeout)
96 try:
97 while True:
98 data = sock.recv(socket_read_size)
99 # an empty string indicates the server shutdown the socket
100 if isinstance(data, bytes) and len(data) == 0:
101 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
102 buf.write(data)
103 data_length = len(data)
104 marker += data_length
106 if length is not None and length > marker:
107 continue
108 return True
109 except socket.timeout:
110 if raise_on_timeout:
111 raise TimeoutError("Timeout reading from socket")
112 return False
113 except NONBLOCKING_EXCEPTIONS as ex:
114 # if we're in nonblocking mode and the recv raises a
115 # blocking error, simply return False indicating that
116 # there's no data to be read. otherwise raise the
117 # original exception.
118 allowed = NONBLOCKING_EXCEPTION_ERROR_NUMBERS.get(ex.__class__, -1)
119 if ex.errno == allowed:
120 if not raise_on_timeout:
121 return False
122 if timeout == 0:
123 raise TimeoutError("Timeout reading from socket")
124 raise ConnectionError(f"Error while reading from socket: {ex.args}")
125 except ValueError:
126 # The buffer was closed by another thread while this read was in flight.
127 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
128 finally:
129 try:
130 buf.seek(current_pos)
131 except ValueError:
132 # Closed by another thread while recv was blocked. Whatever the body
133 # of the read raised is the outcome to report, so this stays quiet
134 # rather than replacing it - the next read raises ConnectionError.
135 pass
136 if custom_timeout:
137 try:
138 sock.settimeout(self.socket_timeout)
139 except OSError:
140 # Same window as the seek above: the close that dropped the
141 # buffer closed the socket too, so there is nothing left to
142 # restore the timeout on. Staying quiet keeps the outcome the
143 # body reported instead of replacing it with EBADF.
144 pass
146 def can_read(self, timeout: float = 0) -> bool:
147 return bool(self.unread_bytes()) or self._read_from_socket(
148 timeout=timeout, raise_on_timeout=False
149 )
151 def read(self, length: int, timeout: Union[float, object] = SENTINEL) -> bytes:
152 length = length + 2 # make sure to read the \r\n terminator
153 buf = self._live_buffer()
154 try:
155 # BufferIO will return less than requested if buffer is short
156 data = buf.read(length)
157 missing = length - len(data)
158 if missing:
159 # fill up the buffer and read the remainder
160 self._read_from_socket(length=missing, timeout=timeout)
161 data += buf.read(missing)
162 except ValueError:
163 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
164 return data[:-2]
166 def readline(self, timeout: Union[float, object] = SENTINEL) -> bytes:
167 buf = self._live_buffer()
168 try:
169 data = buf.readline()
170 while not data.endswith(SYM_CRLF):
171 # there's more data in the socket that we need
172 self._read_from_socket(timeout=timeout)
173 data += buf.readline()
174 except ValueError:
175 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
177 return data[:-2]
179 def get_pos(self) -> int:
180 """
181 Get current read position
182 """
183 try:
184 return self._live_buffer().tell()
185 except ValueError:
186 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
188 def rewind(self, pos: int) -> None:
189 """
190 Rewind the buffer to a specific position, to re-start reading
191 """
192 buffer = self._buffer
194 # Best effort: the caller is unwinding a read that already failed, and a
195 # buffer closed by another thread has nothing to rewind. Raising here would
196 # replace the exception the caller is propagating.
197 if buffer is None:
198 return
200 try:
201 buffer.seek(pos)
202 except ValueError:
203 pass
205 def purge(self) -> None:
206 """
207 After a successful read, purge the read part of buffer
208 """
209 try:
210 self._purge()
211 except (ConnectionError, ValueError):
212 # Closed by another thread while the response was being read. The
213 # response is already parsed, so there is nothing to report and nothing
214 # left to purge.
215 return
217 def _purge(self) -> None:
218 unread = self.unread_bytes()
220 # Only if we have read all of the buffer do we truncate, to
221 # reduce the amount of memory thrashing. This heuristic
222 # can be changed or removed later.
223 if unread > 0:
224 return
226 # Bind the buffer once: another thread's ``close()`` can drop
227 # ``self._buffer`` to None between the read above and the truncate below,
228 # and the resulting AttributeError would escape ``purge()``'s best effort
229 # wrapper. A local reference to an already closed buffer raises
230 # ValueError, which ``purge()`` handles.
231 buffer = self._live_buffer()
233 if unread > 0:
234 # move unread data to the front
235 view = buffer.getbuffer()
236 view[:unread] = view[-unread:]
237 buffer.truncate(unread)
238 buffer.seek(0)
240 def close(self) -> None:
241 try:
242 self._buffer.close()
243 except Exception:
244 # issue #633 suggests the purge/close somehow raised a
245 # BadFileDescriptor error. Perhaps the client ran out of
246 # memory or something else? It's probably OK to ignore
247 # any error being raised from purge/close since we're
248 # removing the reference to the instance below.
249 pass
250 self._buffer = None
251 self._sock = None