PYTHON-4636 - Avoid blocking I/O calls in async code paths #1870
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -16,15 +16,21 @@ | |
| from __future__ import annotations | ||
|
|
||
| import asyncio | ||
| import errno | ||
| import socket | ||
| import struct | ||
| import sys | ||
| import time | ||
| from asyncio import AbstractEventLoop, Future | ||
| from typing import ( | ||
| TYPE_CHECKING, | ||
| Optional, | ||
| Union, | ||
| ) | ||
|
|
||
| from pymongo import _csot, ssl_support | ||
| from pymongo.errors import _OperationCancelled | ||
| from pymongo.socket_checker import _errno_from_exception | ||
|
|
||
| try: | ||
| from ssl import SSLError, SSLSocket | ||
|
|
@@ -51,6 +57,10 @@ | |
| BLOCKING_IO_WRITE_ERROR, | ||
| ) | ||
|
|
||
| if TYPE_CHECKING: | ||
| from pymongo.asynchronous.pool import AsyncConnection | ||
| from pymongo.synchronous.pool import Connection | ||
|
|
||
| _UNPACK_HEADER = struct.Struct("<iiii").unpack | ||
| _UNPACK_COMPRESSION_HEADER = struct.Struct("<iiB").unpack | ||
| _POLL_TIMEOUT = 0.5 | ||
|
|
@@ -111,6 +121,46 @@ def _is_ready(fut: Future) -> None: | |
| loop.add_reader(fd, _is_ready, fut) | ||
| loop.add_writer(fd, _is_ready, fut) | ||
| await fut | ||
|
|
||
| async def _async_receive_ssl( | ||
| conn: _sslConn, length: int, loop: AbstractEventLoop | ||
| ) -> memoryview: | ||
| mv = memoryview(bytearray(length)) | ||
| fd = conn.fileno() | ||
| read = 0 | ||
|
|
||
| def _is_ready(fut: Future) -> None: | ||
| loop.remove_writer(fd) | ||
| loop.remove_reader(fd) | ||
| if fut.done(): | ||
| return | ||
| fut.set_result(None) | ||
|
|
||
| while read < length: | ||
| try: | ||
| read += conn.recv_into(mv[read:]) | ||
| if read == 0: | ||
| raise OSError("connection closed") | ||
| except BLOCKING_IO_ERRORS as exc: | ||
| fd = conn.fileno() | ||
| # Check for closed socket. | ||
| if fd == -1: | ||
| raise SSLError("Underlying socket has been closed") from None | ||
| if isinstance(exc, BLOCKING_IO_READ_ERROR): | ||
| fut = loop.create_future() | ||
| loop.add_reader(fd, _is_ready, fut) | ||
| await fut | ||
| if isinstance(exc, BLOCKING_IO_WRITE_ERROR): | ||
| fut = loop.create_future() | ||
| loop.add_writer(fd, _is_ready, fut) | ||
| await fut | ||
| if _HAVE_PYOPENSSL and isinstance(exc, BLOCKING_IO_LOOKUP_ERROR): | ||
| fut = loop.create_future() | ||
| loop.add_reader(fd, _is_ready, fut) | ||
| loop.add_writer(fd, _is_ready, fut) | ||
| await fut | ||
| return mv | ||
|
|
||
| else: | ||
| # The default Windows asyncio event loop does not support loop.add_reader/add_writer: | ||
| # https://docs.python.org/3/library/asyncio-platforms.html#asyncio-platform-support | ||
|
|
@@ -128,6 +178,136 @@ async def _async_sendall_ssl( | |
| sent = 0 | ||
| total_sent += sent | ||
|
|
||
| async def _async_receive_ssl( | ||
| conn: _sslConn, length: int, dummy: AbstractEventLoop | ||
| ) -> memoryview: | ||
| mv = memoryview(bytearray(length)) | ||
| total_read = 0 | ||
| while total_read < length: | ||
| try: | ||
| read = conn.recv_into(mv[total_read:]) | ||
| except BLOCKING_IO_ERRORS: | ||
| await asyncio.sleep(0.5) | ||
|
||
| read = 0 | ||
| total_read += read | ||
| return mv | ||
|
|
||
|
|
||
| def sendall(sock: Union[socket.socket, _sslConn], buf: bytes) -> None: | ||
| sock.sendall(buf) | ||
|
|
||
|
|
||
| async def _poll_cancellation(conn: AsyncConnection) -> None: | ||
| while True: | ||
| if conn.cancel_context.cancelled: | ||
| return | ||
|
|
||
| await asyncio.sleep(_POLL_TIMEOUT) | ||
|
|
||
|
|
||
| async def async_receive_data( | ||
| conn: AsyncConnection, length: int, deadline: Optional[float] | ||
| ) -> memoryview: | ||
| sock = conn.conn | ||
| sock_timeout = sock.gettimeout() | ||
| timeout: Optional[Union[float, int]] | ||
| if deadline: | ||
| # When the timeout has expired perform one final check to | ||
| # see if the socket is readable. This helps avoid spurious | ||
| # timeouts on AWS Lambda and other FaaS environments. | ||
| timeout = max(deadline - time.monotonic(), 0) | ||
| else: | ||
| timeout = sock_timeout | ||
|
|
||
| sock.settimeout(0.0) | ||
| loop = asyncio.get_event_loop() | ||
| cancellation_task = asyncio.create_task(_poll_cancellation(conn)) | ||
| try: | ||
| if _HAVE_SSL and isinstance(sock, (SSLSocket, _sslConn)): | ||
| read_task = asyncio.create_task(_async_receive_ssl(sock, length, loop)) # type: ignore[arg-type] | ||
| else: | ||
| read_task = asyncio.create_task(_async_receive(sock, length, loop)) # type: ignore[arg-type] | ||
| tasks = [read_task, cancellation_task] | ||
| result = await asyncio.wait(tasks, timeout=timeout, return_when=asyncio.FIRST_COMPLETED) | ||
| if len(result[1]) == 2: | ||
| raise socket.timeout("timed out") | ||
ShaneHarvey marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| finished = next(iter(result[0])) | ||
| next(iter(result[1])).cancel() | ||
| if finished == read_task: | ||
| return finished.result() # type: ignore[return-value] | ||
|
||
| else: | ||
| raise _OperationCancelled("operation cancelled") | ||
| finally: | ||
| sock.settimeout(sock_timeout) | ||
|
|
||
|
|
||
| async def _async_receive(conn: socket.socket, length: int, loop: AbstractEventLoop) -> memoryview: | ||
| mv = memoryview(bytearray(length)) | ||
| bytes_read = 0 | ||
| while bytes_read < length: | ||
| chunk_length = await loop.sock_recv_into(conn, mv[bytes_read:]) | ||
| if chunk_length == 0: | ||
| raise OSError("connection closed") | ||
| bytes_read += chunk_length | ||
| return mv | ||
|
|
||
|
|
||
| # Sync version: | ||
| def wait_for_read(conn: Connection, deadline: Optional[float]) -> None: | ||
| """Block until at least one byte is read, or a timeout, or a cancel.""" | ||
| sock = conn.conn | ||
| timed_out = False | ||
| # Check if the connection's socket has been manually closed | ||
| if sock.fileno() == -1: | ||
| return | ||
| while True: | ||
| # SSLSocket can have buffered data which won't be caught by select. | ||
| if hasattr(sock, "pending") and sock.pending() > 0: | ||
| readable = True | ||
| else: | ||
| # Wait up to 500ms for the socket to become readable and then | ||
| # check for cancellation. | ||
| if deadline: | ||
| remaining = deadline - time.monotonic() | ||
| # When the timeout has expired perform one final check to | ||
| # see if the socket is readable. This helps avoid spurious | ||
| # timeouts on AWS Lambda and other FaaS environments. | ||
| if remaining <= 0: | ||
| timed_out = True | ||
| timeout = max(min(remaining, _POLL_TIMEOUT), 0) | ||
| else: | ||
| timeout = _POLL_TIMEOUT | ||
| readable = conn.socket_checker.select(sock, read=True, timeout=timeout) | ||
| if conn.cancel_context.cancelled: | ||
| raise _OperationCancelled("operation cancelled") | ||
| if readable: | ||
| return | ||
| if timed_out: | ||
| raise socket.timeout("timed out") | ||
|
|
||
|
|
||
| def receive_data(conn: Connection, length: int, deadline: Optional[float]) -> memoryview: | ||
| buf = bytearray(length) | ||
| mv = memoryview(buf) | ||
| bytes_read = 0 | ||
| while bytes_read < length: | ||
| try: | ||
| wait_for_read(conn, deadline) | ||
| # CSOT: Update timeout. When the timeout has expired perform one | ||
| # final non-blocking recv. This helps avoid spurious timeouts when | ||
| # the response is actually already buffered on the client. | ||
| if _csot.get_timeout() and deadline is not None: | ||
| conn.set_conn_timeout(max(deadline - time.monotonic(), 0)) | ||
| chunk_length = conn.conn.recv_into(mv[bytes_read:]) | ||
| except BLOCKING_IO_ERRORS: | ||
| raise socket.timeout("timed out") from None | ||
| except OSError as exc: | ||
| if _errno_from_exception(exc) == errno.EINTR: | ||
| continue | ||
| raise | ||
| if chunk_length == 0: | ||
| raise OSError("connection closed") | ||
|
|
||
| bytes_read += chunk_length | ||
|
|
||
| return mv | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This should be something like: