|
1 | 1 | import logging |
| 2 | +import sys |
2 | 3 | from abc import ABC, abstractmethod |
3 | 4 | from asyncio import IncompleteReadError, StreamReader |
4 | 5 | from typing import Awaitable, Callable, List, Optional, Protocol, Union |
5 | 6 |
|
| 7 | +if sys.version_info >= (3, 11, 3): |
| 8 | + from asyncio import timeout as async_timeout |
| 9 | +else: |
| 10 | + from async_timeout import timeout as async_timeout |
| 11 | + |
6 | 12 | from redis.maint_notifications import ( |
7 | 13 | MaintenanceNotification, |
8 | 14 | NodeFailedOverNotification, |
|
13 | 19 | OSSNodeMigratedNotification, |
14 | 20 | OSSNodeMigratingNotification, |
15 | 21 | ) |
16 | | -from redis.utils import deprecated_function, safe_str |
17 | 22 |
|
18 | 23 | from ..exceptions import ( |
19 | 24 | AskError, |
|
36 | 41 | TryAgainError, |
37 | 42 | ) |
38 | 43 | from ..typing import EncodableT |
| 44 | +from ..utils import SENTINEL, deprecated_function, safe_str |
39 | 45 | from .encoders import Encoder |
40 | 46 | from .socket import SERVER_CLOSED_CONNECTION_ERROR, SocketBuffer |
41 | 47 |
|
@@ -183,7 +189,10 @@ async def can_read(self) -> bool: |
183 | 189 | pass |
184 | 190 |
|
185 | 191 | async def read_response( |
186 | | - self, disable_decoding: bool = False |
| 192 | + self, |
| 193 | + disable_decoding: bool = False, |
| 194 | + push_request: bool = False, |
| 195 | + timeout: Union[float, object] = SENTINEL, |
187 | 196 | ) -> Union[EncodableT, ResponseError, None, List[EncodableT]]: |
188 | 197 | raise NotImplementedError() |
189 | 198 |
|
@@ -545,40 +554,65 @@ async def can_read(self) -> bool: |
545 | 554 | # parser and fail loudly if the private buffer API changes. |
546 | 555 | return bool(self._stream._buffer) or self._stream.at_eof() |
547 | 556 |
|
548 | | - async def _read(self, length: int) -> bytes: |
| 557 | + async def _read_from_stream( |
| 558 | + self, timeout: Union[float, object] = SENTINEL, max_bytes: int = 0 |
| 559 | + ) -> bytes: |
| 560 | + """ |
| 561 | + Read the next chunk from the underlying stream with an optional |
| 562 | + per-read timeout. This mirrors the sync client's per-recv timeout |
| 563 | + semantics: each individual socket read gets its own timeout window. |
| 564 | +
|
| 565 | + ``max_bytes`` limits how many bytes may be returned. When 0, the |
| 566 | + parser's ``_read_size`` is used. |
| 567 | + """ |
| 568 | + stream = self._stream |
| 569 | + if stream is None: |
| 570 | + raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) |
| 571 | + size = max_bytes if max_bytes > 0 else self._read_size |
| 572 | + if timeout is not SENTINEL and isinstance(timeout, (int, float)): |
| 573 | + async with async_timeout(timeout): |
| 574 | + return await stream.read(size) |
| 575 | + return await stream.read(size) |
| 576 | + |
| 577 | + async def _read( |
| 578 | + self, length: int, timeout: Union[float, object] = SENTINEL |
| 579 | + ) -> bytes: |
549 | 580 | """ |
550 | 581 | Read `length` bytes of data. These are assumed to be followed |
551 | 582 | by a '\r\n' terminator which is subsequently discarded. |
552 | 583 | """ |
553 | 584 | want = length + 2 |
554 | 585 | end = self._pos + want |
555 | | - if len(self._buffer) >= end: |
556 | | - result = self._buffer[self._pos : end - 2] |
557 | | - else: |
558 | | - tail = self._buffer[self._pos :] |
| 586 | + while len(self._buffer) < end: |
| 587 | + need = end - len(self._buffer) |
559 | 588 | try: |
560 | | - data = await self._stream.readexactly(want - len(tail)) |
| 589 | + chunk = await self._read_from_stream( |
| 590 | + timeout=timeout, max_bytes=min(need, self._read_size) |
| 591 | + ) |
561 | 592 | except IncompleteReadError as error: |
562 | 593 | raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from error |
563 | | - result = (tail + data)[:-2] |
564 | | - self._chunks.append(data) |
565 | | - self._pos += want |
| 594 | + if not chunk: |
| 595 | + raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) |
| 596 | + self._buffer += chunk |
| 597 | + result = self._buffer[self._pos : end - 2] |
| 598 | + self._pos = end |
566 | 599 | return result |
567 | 600 |
|
568 | | - async def _readline(self) -> bytes: |
| 601 | + async def _readline(self, timeout: Union[float, object] = SENTINEL) -> bytes: |
569 | 602 | """ |
570 | 603 | read an unknown number of bytes up to the next '\r\n' |
571 | 604 | line separator, which is discarded. |
572 | 605 | """ |
573 | | - found = self._buffer.find(b"\r\n", self._pos) |
574 | | - if found >= 0: |
575 | | - result = self._buffer[self._pos : found] |
576 | | - else: |
577 | | - tail = self._buffer[self._pos :] |
578 | | - data = await self._stream.readline() |
579 | | - if not data.endswith(b"\r\n"): |
| 606 | + while True: |
| 607 | + found = self._buffer.find(b"\r\n", self._pos) |
| 608 | + if found >= 0: |
| 609 | + result = self._buffer[self._pos : found] |
| 610 | + self._pos = found + 2 |
| 611 | + return result |
| 612 | + try: |
| 613 | + chunk = await self._read_from_stream(timeout=timeout) |
| 614 | + except IncompleteReadError as error: |
| 615 | + raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from error |
| 616 | + if not chunk: |
580 | 617 | raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) |
581 | | - result = (tail + data)[:-2] |
582 | | - self._chunks.append(data) |
583 | | - self._pos += len(result) + 2 |
584 | | - return result |
| 618 | + self._buffer += chunk |
0 commit comments