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 select
2import selectors
3import socket
4from logging import getLogger
5from typing import Callable, List, Optional, TypedDict, Union
7from ..exceptions import ConnectionError, InvalidResponse, RedisError, TimeoutError
8from ..typing import EncodableT
9from ..utils import HIREDIS_AVAILABLE, SENTINEL, deprecated_function
10from .base import (
11 AsyncBaseParser,
12 AsyncPushNotificationsParser,
13 BaseParser,
14 PushNotificationsParser,
15)
16from .socket import (
17 NONBLOCKING_EXCEPTION_ERROR_NUMBERS,
18 NONBLOCKING_EXCEPTIONS,
19 SERVER_CLOSED_CONNECTION_ERROR,
20)
22# Used to signal that hiredis-py does not have enough data to parse.
23# Using `False` or `None` is not reliable, given that the parser can
24# return `False` or `None` for legitimate reasons from RESP payloads.
25NOT_ENOUGH_DATA = object()
27# select.poll() is unavailable on Windows; fall back to selectors there.
28_HAS_POLL = hasattr(select, "poll")
30# POLLRDHUP (Linux) reports a peer half-close (FIN) that POLLHUP does not cover.
31# It is absent on macOS/Windows, where 0 turns the masks that use it into no-ops.
32_POLLRDHUP = getattr(select, "POLLRDHUP", 0)
35def _socket_can_read(sock, timeout: float) -> bool:
36 # SSL sockets can have decrypted bytes buffered above the OS socket layer.
37 if hasattr(sock, "pending") and sock.pending():
38 return True
39 # timeout=0 must be a non-blocking readiness check only; both branches
40 # below are non-destructive and have no FD_SETSIZE limit (select.select
41 # raises ValueError for fds >= 1024).
42 if _HAS_POLL:
43 # Prefer poll() over selectors.DefaultSelector: epoll/kqueue selectors
44 # allocate a file descriptor per check and so fail with EMFILE under
45 # fd exhaustion - the very condition that pushes sockets onto high
46 # fds. poll() allocates nothing.
47 poller = select.poll()
48 poller.register(sock, select.POLLIN)
49 # poll() takes milliseconds (None blocks forever). POLLHUP/POLLERR/
50 # POLLNVAL are always reported regardless of the registered mask, so
51 # closed or errored sockets still count as readable, like select().
52 poll_timeout = None if timeout is None else timeout * 1000
53 return bool(poller.poll(poll_timeout))
54 with selectors.DefaultSelector() as selector:
55 selector.register(sock, selectors.EVENT_READ)
56 return bool(selector.select(timeout))
59def _socket_is_closed(sock) -> bool:
60 # A server-closed socket reads as ready (it yields EOF), so readiness alone
61 # cannot tell it apart from a socket holding pending data, and both checks
62 # here are non-destructive so pending push messages (e.g. cache
63 # invalidations) are left intact to be processed. Without poll() the two
64 # states are indistinguishable, so report not-closed.
65 if not _HAS_POLL:
66 return False
67 # Decrypted TLS bytes buffered above the OS socket layer must be processed
68 # before the connection can be treated as closed, like kernel-level data.
69 if hasattr(sock, "pending") and sock.pending():
70 return False
71 poller = select.poll()
72 poller.register(sock, select.POLLIN | _POLLRDHUP)
73 events = poller.poll(0)
74 if not events:
75 return False
76 _, revents = events[0]
77 # A readable socket holds either data or EOF, and the poll flags alone
78 # cannot always tell which: POLLHUP/POLLRDHUP can be reported while unread
79 # data is still buffered (the peer sent data, then closed), and a drained
80 # closed socket reports plain POLLIN where POLLRDHUP is unavailable
81 # (PyPy). A non-destructive MSG_PEEK settles it: EOF peeks as b"".
82 try:
83 return sock.recv(1, socket.MSG_PEEK) == b""
84 except NONBLOCKING_EXCEPTIONS:
85 # Transient would-block: data may still arrive, keep the connection.
86 return False
87 except ValueError:
88 # SSL sockets do not support recv() flags; their buffered plaintext is
89 # covered by the pending() check above, so fall back to the poll
90 # flags. POLLRDHUP must be in the register mask or poll() won't
91 # report it: on Linux a graceful FIN reports POLLIN|POLLRDHUP and
92 # never POLLHUP. macOS/Windows set POLLHUP instead (_POLLRDHUP is 0).
93 closed_flags = select.POLLHUP | select.POLLERR | select.POLLNVAL | _POLLRDHUP
94 return bool(revents & closed_flags)
95 except OSError:
96 # The socket is errored (POLLERR/POLLNVAL); nothing left to read.
97 return True
100class _HiredisReaderArgs(TypedDict, total=False):
101 protocolError: Callable[[str], Exception]
102 replyError: Callable[[str], Exception]
103 encoding: Optional[str]
104 errors: Optional[str]
107class _HiredisParser(BaseParser, PushNotificationsParser):
108 "Parser class for connections using Hiredis"
110 def __init__(self, socket_read_size):
111 if not HIREDIS_AVAILABLE:
112 raise RedisError("Hiredis is not installed")
113 self.socket_read_size = socket_read_size
114 self._buffer = bytearray(socket_read_size)
115 self.pubsub_push_handler_func = self.handle_pubsub_push_response
116 self.node_moving_push_handler_func = None
117 self.maintenance_push_handler_func = None
118 self.oss_cluster_maint_push_handler_func = None
119 self.invalidation_push_handler_func = None
120 self._hiredis_PushNotificationType = None
122 def __del__(self):
123 try:
124 self.on_disconnect()
125 except Exception:
126 pass
128 def handle_pubsub_push_response(self, response):
129 logger = getLogger("push_response")
130 logger.debug("Push response: " + str(response))
131 return response
133 def on_connect(self, connection, **kwargs):
134 import hiredis
136 self._sock = connection._sock
137 self._socket_timeout = connection.socket_timeout
138 kwargs = {
139 "protocolError": InvalidResponse,
140 "replyError": self.parse_error,
141 "errors": connection.encoder.encoding_errors,
142 "notEnoughData": NOT_ENOUGH_DATA,
143 }
145 if connection.encoder.decode_responses:
146 kwargs["encoding"] = connection.encoder.encoding
147 self._reader = hiredis.Reader(**kwargs)
149 try:
150 self._hiredis_PushNotificationType = hiredis.PushNotification
151 except AttributeError:
152 # hiredis < 3.2
153 self._hiredis_PushNotificationType = None
155 def on_disconnect(self):
156 self._sock = None
157 self._reader = None
159 def can_read(self, timeout: float = 0) -> bool:
160 # TODO: Rename this API; it detects pending data or dirty/closed
161 # connection state, not only whether application data can be read.
162 if not self._reader:
163 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
165 if self._reader.has_data():
166 return True
167 if not _socket_can_read(self._sock, timeout):
168 return False
169 # the socket reports readable but the reader has no buffered data. a
170 # server-closed socket also reads as ready (it yields EOF), so tell the
171 # two apart with a non-destructive poll: a peer-closed socket must not be
172 # reused, while a readable-but-open socket may just hold a pending push.
173 # this mirrors how the pure-Python parser (recv -> b"") and the async
174 # parser (StreamReader.at_eof()) already signal a closed connection.
175 if _socket_is_closed(self._sock):
176 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
177 return True
179 def read_from_socket(self, timeout=SENTINEL, raise_on_timeout=True):
180 sock = self._sock
181 reader = self._reader
182 # Another thread may disconnect this connection while we are here (e.g.
183 # a shared client closed via `with redis:`); on_disconnect() sets both
184 # _sock and _reader to None. Bind them locally and fail with a
185 # descriptive, retryable ConnectionError instead of an AttributeError.
186 if sock is None or reader is None:
187 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
188 custom_timeout = timeout is not SENTINEL
189 try:
190 if custom_timeout:
191 sock.settimeout(timeout)
192 bufflen = sock.recv_into(self._buffer)
193 if bufflen == 0:
194 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
195 reader.feed(self._buffer, 0, bufflen)
196 # data was read from the socket and added to the buffer.
197 # return True to indicate that data was read.
198 return True
199 except socket.timeout:
200 if raise_on_timeout:
201 raise TimeoutError("Timeout reading from socket")
202 return False
203 except NONBLOCKING_EXCEPTIONS as ex:
204 # if we're in nonblocking mode and the recv raises a
205 # blocking error, simply return False indicating that
206 # there's no data to be read. otherwise raise the
207 # original exception.
208 allowed = NONBLOCKING_EXCEPTION_ERROR_NUMBERS.get(ex.__class__, -1)
209 if ex.errno == allowed:
210 if not raise_on_timeout:
211 return False
212 if timeout == 0:
213 raise TimeoutError("Timeout reading from socket")
214 raise ConnectionError(f"Error while reading from socket: {ex.args}")
215 finally:
216 if custom_timeout:
217 sock.settimeout(self._socket_timeout)
219 def read_response(
220 self,
221 disable_decoding=False,
222 push_request=False,
223 timeout: Union[float, object] = SENTINEL,
224 ):
225 # Bind the reader locally so a concurrent disconnect that clears
226 # self._reader can't turn a later .gets() into an AttributeError;
227 # re-checking the attribute each time would still race.
228 reader = self._reader
229 if reader is None:
230 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
232 if disable_decoding:
233 response = reader.gets(False)
234 else:
235 response = reader.gets()
237 while response is NOT_ENOUGH_DATA:
238 self.read_from_socket(timeout=timeout)
239 if disable_decoding:
240 response = reader.gets(False)
241 else:
242 response = reader.gets()
243 # if the response is a ConnectionError or the response is a list and
244 # the first item is a ConnectionError, raise it as something bad
245 # happened
246 if isinstance(response, ConnectionError):
247 raise response
248 elif self._hiredis_PushNotificationType is not None and isinstance(
249 response, self._hiredis_PushNotificationType
250 ):
251 response = self.handle_push_response(response)
252 if push_request:
253 return response
254 return self.read_response(
255 disable_decoding=disable_decoding,
256 push_request=push_request,
257 timeout=timeout,
258 )
260 elif (
261 isinstance(response, list)
262 and response
263 and isinstance(response[0], ConnectionError)
264 ):
265 raise response[0]
266 return response
269class _AsyncHiredisParser(AsyncBaseParser, AsyncPushNotificationsParser):
270 """Async implementation of parser class for connections using Hiredis"""
272 __slots__ = ("_reader",)
274 def __init__(self, socket_read_size: int):
275 if not HIREDIS_AVAILABLE:
276 raise RedisError("Hiredis is not available.")
277 super().__init__(socket_read_size=socket_read_size)
278 self._reader = None
279 self.pubsub_push_handler_func = self.handle_pubsub_push_response
280 self.invalidation_push_handler_func = None
281 self._hiredis_PushNotificationType = None
283 async def handle_pubsub_push_response(self, response):
284 logger = getLogger("push_response")
285 logger.debug("Push response: " + str(response))
286 return response
288 def on_connect(self, connection):
289 import hiredis
291 self._stream = connection._reader
292 kwargs: _HiredisReaderArgs = {
293 "protocolError": InvalidResponse,
294 "replyError": self.parse_error,
295 "notEnoughData": NOT_ENOUGH_DATA,
296 }
297 if connection.encoder.decode_responses:
298 kwargs["encoding"] = connection.encoder.encoding
299 kwargs["errors"] = connection.encoder.encoding_errors
301 self._reader = hiredis.Reader(**kwargs)
302 self._connected = True
304 try:
305 self._hiredis_PushNotificationType = getattr(
306 hiredis, "PushNotification", None
307 )
308 except AttributeError:
309 # hiredis < 3.2
310 self._hiredis_PushNotificationType = None
312 def on_disconnect(self):
313 self._connected = False
315 @deprecated_function(
316 version="8.0.0", reason="Use can_read() instead", name="can_read_destructive"
317 )
318 async def can_read_destructive(self) -> bool:
319 return await self.can_read()
321 async def can_read(self) -> bool:
322 # TODO: Rename this API; it detects pending data or dirty/closed
323 # connection state, not only whether application data can be read.
324 if not self._connected:
325 raise OSError("Buffer is closed.")
326 # EOF means the connection is closed and not safe to reuse.
327 if self._reader.has_data() or self._stream.at_eof():
328 return True
329 # asyncio.StreamReader has no public non-destructive API for checking
330 # buffered bytes. Preserve dirty-connection detection for hiredis; tests
331 # with a real StreamReader guard this private buffer API in CI.
332 return bool(self._stream._buffer)
334 async def read_from_socket(self):
335 buffer = await self._stream.read(self._read_size)
336 if not buffer or not isinstance(buffer, bytes):
337 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
338 self._reader.feed(buffer)
339 # data was read from the socket and added to the buffer.
340 # return True to indicate that data was read.
341 return True
343 async def read_response(
344 self, disable_decoding: bool = False, push_request: bool = False
345 ) -> Union[EncodableT, List[EncodableT]]:
346 # If `on_disconnect()` has been called, prohibit any more reads
347 # even if they could happen because data might be present.
348 # We still allow reads in progress to finish
349 if not self._connected:
350 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
352 if disable_decoding:
353 response = self._reader.gets(False)
354 else:
355 response = self._reader.gets()
357 while response is NOT_ENOUGH_DATA:
358 await self.read_from_socket()
359 if disable_decoding:
360 response = self._reader.gets(False)
361 else:
362 response = self._reader.gets()
364 # if the response is a ConnectionError or the response is a list and
365 # the first item is a ConnectionError, raise it as something bad
366 # happened
367 if isinstance(response, ConnectionError):
368 raise response
369 elif self._hiredis_PushNotificationType is not None and isinstance(
370 response, self._hiredis_PushNotificationType
371 ):
372 response = await self.handle_push_response(response)
373 if not push_request:
374 return await self.read_response(
375 disable_decoding=disable_decoding, push_request=push_request
376 )
377 else:
378 return response
379 elif (
380 isinstance(response, list)
381 and response
382 and isinstance(response[0], ConnectionError)
383 ):
384 raise response[0]
385 return response